use pytest-mock for the VariantFetcherProxy

Author: Hoeze

kipoiseq/extractors/vcf_matching.py

@@ -110,21 +110,6 @@ def __iter__(self) -> Iterator[Variant]:
         yield from self.variants
 
 
-class VariantFetcherProxy(VariantFetcher):
-
-    def __init__(self, variant_fetcher: VariantFetcher):
-        self.variant_fetcher = variant_fetcher
-
-    def fetch_variants(self, interval: Union[Interval, Iterable[Interval]]) -> Iterator[Variant]:
-        yield from self.variant_fetcher.fetch_variants(interval)
-
-    def batch_iter(self, batch_size=10000) -> Iterator[List[Variant]]:
-        yield from self.variant_fetcher.batch_iter(batch_size)
-
-    def __iter__(self) -> Iterator[Variant]:
-        yield from self.variant_fetcher
-
-
 class BaseVariantMatcher:
     """
     Base variant intervals matcher
@@ -168,16 +153,10 @@ def _read_variants(
             variants=None,
             variant_fetcher=None,
             vcf_lazy: bool = True,
-    ):
+    ) -> VariantFetcher:
         if vcf_file is not None:
             from kipoiseq.extractors import MultiSampleVCF
-            vcf = MultiSampleVCF(vcf_file, lazy=vcf_lazy)
-
-            if os.environ.get('PYTEST_RUNNING', '') == 'true':
-                # Ensure that none of the methods actually uses MultiSampleVCF methods by accident
-                vcf = VariantFetcherProxy(vcf)
-
-            return vcf
+            return MultiSampleVCF(vcf_file, lazy=vcf_lazy)
         elif variant_fetcher is not None:
             assert isinstance(variant_fetcher, VariantFetcher), \
                 "Wrong type of variant fetcher: %s" % type(variant_fetcher)

tests/extractors/test_vcf_matching.py

@@ -1,11 +1,12 @@
+from typing import Union, Iterable, Iterator, List
 import pytest
 from conftest import vcf_file, gtf_file, example_intervals_bed
 import pyranges
 from kipoiseq.dataclasses import Interval, Variant
 from kipoiseq.extractors.vcf import MultiSampleVCF
 from kipoiseq.extractors.vcf_matching import variants_to_pyranges, \
     pyranges_to_intervals, intervals_to_pyranges, BaseVariantMatcher, \
-    SingleVariantMatcher, MultiVariantsMatcher
+    SingleVariantMatcher, MultiVariantsMatcher, VariantFetcher
 
 intervals = [
     Interval('chr1', 1, 10, strand='+'),
@@ -26,6 +27,36 @@
 )
 
 
+class VariantFetcherProxy(VariantFetcher):
+
+    def __init__(self, variant_fetcher: VariantFetcher):
+        self.variant_fetcher = variant_fetcher
+
+    def fetch_variants(self, interval: Union[Interval, Iterable[Interval]]) -> Iterator[Variant]:
+        yield from self.variant_fetcher.fetch_variants(interval)
+
+    def batch_iter(self, batch_size=10000) -> Iterator[List[Variant]]:
+        yield from self.variant_fetcher.batch_iter(batch_size)
+
+    def __iter__(self) -> Iterator[Variant]:
+        yield from self.variant_fetcher
+
+
+# make sure that kipoiseq only uses the VariantFetcher API
+read_variants_fn = BaseVariantMatcher._read_variants
+
+
+@staticmethod
+def proxy_fn(*args, **kwargs):
+    vf = VariantFetcherProxy(
+        read_variants_fn(*args, **kwargs)
+    )
+    return vf
+
+
+BaseVariantMatcher._read_variants = proxy_fn
+
+
 def test_variants_to_pyranges():
     vcf = MultiSampleVCF(vcf_file)
     variants = list(vcf)
@@ -140,6 +171,7 @@ def test_SingleVariantMatcher__iter__():
     assert (inters[2], variants[2]) in pairs
     assert len(pairs) == 4
 
+
 def test_MultiVariantMatcher__iter__():
     matcher = MultiVariantsMatcher(vcf_file, intervals=intervals)
     pairs = list(matcher)