fix tests

broadinstitute · Feb 13, 2024 · 03332b0 · 03332b0
1 parent ef4a8bc
commit 03332b0
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/seqr/utils/search/hail_search_utils.py b/seqr/utils/search/hail_search_utils.py
@@ -75,22 +75,23 @@ def get_hail_variants_for_variant_ids(samples, genome_version, parsed_variant_id
 
 
 def hail_variant_lookup(user, variant_id, samples=None, dataset_type=Sample.DATASET_TYPE_VARIANT_CALLS, sample_type=None, **kwargs):
-    is_sv = dataset_type == Sample.DATASET_TYPE_SV_CALLS
+    data_type = dataset_type.replace('_only', '')
+    is_sv = data_type == Sample.DATASET_TYPE_SV_CALLS
     if is_sv:
         if not sample_type:
             from seqr.utils.search.utils import InvalidSearchException
             raise InvalidSearchException('Sample type must be specified to look up a structural variant')
-        dataset_type = f'{dataset_type}_{sample_type}'
+        data_type = f'{data_type}_{sample_type}'
 
     body = {
         'variant_id': variant_id,
-        'data_type': dataset_type.replace('_only', ''),
+        'data_type': data_type,
         **kwargs,
     }
     sample_data = None
     if samples:
         sample_data = _get_sample_data(samples)
-        body['sample_data'] = sample_data.pop(dataset_type)
+        body['sample_data'] = sample_data.pop(data_type)
     variant = _execute_search(body, user, path='lookup', exception_map={404: 'Variant not present in seqr'})
     variants = [variant]
 

diff --git a/seqr/utils/search/hail_search_utils_tests.py b/seqr/utils/search/hail_search_utils_tests.py
@@ -214,28 +214,29 @@ def test_get_variant_query_gene_counts(self):
     def test_variant_lookup(self):
         responses.add(responses.POST, f'{MOCK_HOST}:5000/lookup', status=200, json=VARIANT_LOOKUP_VARIANT)
         variant = variant_lookup(self.user, '1-10439-AC-A', genome_version='37', foo='bar')
-        self.assertDictEqual(variant, VARIANT_LOOKUP_VARIANT)
+        self.assertListEqual(variant, [VARIANT_LOOKUP_VARIANT])
         self._test_minimal_search_call(expected_search_body={
-            'variant_id': ['1', 10439, 'AC', 'A'], 'genome_version': 'GRCh37', 'foo': 'bar',
+            'variant_id': ['1', 10439, 'AC', 'A'], 'genome_version': 'GRCh37', 'foo': 'bar', 'data_type': 'SNV_INDEL',
         })
 
         variant_lookup(self.user, '1-10439-AC-A', genome_version='37', families=self.families)
         self._test_minimal_search_call(expected_search_body={
             'variant_id': ['1', 10439, 'AC', 'A'], 'genome_version': 'GRCh37',
-            'sample_data': ALL_AFFECTED_SAMPLE_DATA['SNV_INDEL'],
+            'sample_data': ALL_AFFECTED_SAMPLE_DATA['SNV_INDEL'], 'data_type': 'SNV_INDEL',
         })
 
         with self.assertRaises(InvalidSearchException) as cm:
             variant_lookup(self.user, 'prefix_123_DEL')
-        self.assertEqual(str(cm.exception), 'Invalid variant prefix_123_DEL')
+        self.assertEqual(str(cm.exception), 'Sample type must be specified to look up a structural variant')
+        # TODO add test for SVs
 
         responses.add(responses.POST, f'{MOCK_HOST}:5000/lookup', status=404)
         with self.assertRaises(HTTPError) as cm:
             variant_lookup(self.user, '1-10439-AC-A')
         self.assertEqual(cm.exception.response.status_code, 404)
         self.assertEqual(str(cm.exception), 'Variant not present in seqr')
         self._test_minimal_search_call(expected_search_body={
-            'variant_id': ['1', 10439, 'AC', 'A'], 'genome_version': 'GRCh38'
+            'variant_id': ['1', 10439, 'AC', 'A'], 'genome_version': 'GRCh38', 'data_type': 'SNV_INDEL',
         })
 
     @responses.activate

diff --git a/seqr/utils/search/search_utils_tests.py b/seqr/utils/search/search_utils_tests.py
@@ -55,20 +55,18 @@ def test_variant_lookup(self, mock_variant_lookup):
         mock_variant_lookup.return_value = VARIANT_LOOKUP_VARIANT
         variant = variant_lookup(self.user, '1-10439-AC-A', genome_version='38')
         self.assertDictEqual(variant, VARIANT_LOOKUP_VARIANT)
-        mock_variant_lookup.assert_called_with(self.user, ('1', 10439, 'AC', 'A'), genome_version='GRCh38')
+        mock_variant_lookup.assert_called_with(self.user, ('1', 10439, 'AC', 'A'), genome_version='GRCh38',
+                                               dataset_type='SNV_INDEL_only')
         cache_key = 'variant_lookup_results__1-10439-AC-A__38__test_user'
         self.assert_cached_results(variant, cache_key=cache_key)
 
         variant = variant_lookup(self.user, '1-10439-AC-A', genome_version='37', families=self.families)
         self.assertDictEqual(variant, VARIANT_LOOKUP_VARIANT)
-        mock_variant_lookup.assert_called_with(self.user, ('1', 10439, 'AC', 'A'), genome_version='GRCh37', samples=mock.ANY)
+        mock_variant_lookup.assert_called_with(self.user, ('1', 10439, 'AC', 'A'), genome_version='GRCh37', samples=mock.ANY,
+                                               dataset_type='SNV_INDEL_only')
         expected_samples = {s for s in self.search_samples if s.guid not in NON_SNP_INDEL_SAMPLES}
         self.assertSetEqual(set(mock_variant_lookup.call_args.kwargs['samples']), expected_samples)
 
-        with self.assertRaises(InvalidSearchException) as cm:
-            variant_lookup(self.user, '100-10439-AC-A')
-        self.assertEqual(str(cm.exception), 'Invalid variant 100-10439-AC-A')
-
         mock_variant_lookup.reset_mock()
         self.set_cache(variant)
         cached_variant = variant_lookup(self.user, '1-10439-AC-A', genome_version='38')