changed genomes to be queried by genome instead of genome sequence

rescript/bv_brc.py

@@ -25,7 +25,7 @@
 
 def get_bv_brc_metadata(
         ids_metadata: Union[qiime2.NumericMetadataColumn,
-                            qiime2.CategoricalMetadataColumn] = None,
+        qiime2.CategoricalMetadataColumn] = None,
         data_type: str = None,
         rql_query: str = None,
         data_field: str = None,
@@ -67,29 +67,48 @@ def get_bv_brc_genomes(
     # Parameter validation
     rql_query = parameter_validation(rql_query=rql_query,
                                      ids=ids,
-                                     data_type="genome_sequence",
+                                     data_type="genome",
                                      data_field=data_field,
                                      metadata=ids_metadata
                                      )
 
-    # Get requests response for genome sequences
-    sequences = download_data(data_type="genome_sequence",
-                              query=rql_query,
-                              accept="application/json",
-                              )
+    response_genomes = download_data(data_type="genome",
+                                     query=rql_query,
+                                     accept="application/json",
+                                     select=["genome_id", "taxon_id"]
+                                     )
 
-    # Convert sequences in JSON to FASTA file
-    genomes = create_genome_fasta(genome_sequences=sequences)
+    # Get genome sequences and create FASTA files
+    genomes = get_genome_sequences(response_genomes=response_genomes)
 
     # Get taxonomy for sequences
-    taxonomy = get_taxonomy(response_sequences=sequences,
+    taxonomy = get_taxonomy(response_sequences=response_genomes,
                             ranks=ranks,
                             rank_propagation=rank_propagation,
                             accession_name="accession")
 
     return genomes, taxonomy
 
 
+def get_genome_sequences(response_genomes):
+    # Extract genome ids from response (list of dicts)
+    genome_ids = set([str(entry['genome_id']) for entry in response_genomes])
+
+    # Fetch the genome sequences for all genome ids
+    genome_sequences = download_data(
+        data_type="genome_sequence",
+        query=f"in(genome_id,({','.join(genome_ids)}))",
+        accept="application/json",
+        select=["accession", "description", "genome_name",
+                "genome_id", "sequence"]
+    )
+
+    # Create FASTA files from sequences
+    genomes = create_genome_fasta(genome_sequences=genome_sequences)
+
+    return genomes
+
+
 def get_bv_brc_genome_features(
         ids_metadata: Union[qiime2.NumericMetadataColumn,
                             qiime2.CategoricalMetadataColumn] = None,

rescript/tests/test_bv_brc.py

@@ -16,7 +16,7 @@
     get_bv_brc_genome_features, parameter_validation, \
     error_handling, download_data, create_genome_fasta, \
     create_taxonomy_entry, get_loci, read_tsv_data_with_dtypes, process_loci, \
-    get_sequences, get_taxonomy, create_taxonomy
+    get_sequences, get_taxonomy, create_taxonomy, get_genome_sequences
 
 
 class TestParameterValidation(TestPluginBase):
@@ -343,20 +343,12 @@ class TestGetBvBrcGenomes(TestPluginBase):
     package = 'rescript.tests'
 
     @patch('rescript.bv_brc.get_taxonomy')
-    @patch('rescript.bv_brc.create_genome_fasta')
+    @patch('rescript.bv_brc.get_genome_sequences')
     @patch('rescript.bv_brc.download_data')
     @patch('rescript.bv_brc.parameter_validation')
     def test_get_bv_brc_genomes(self, mock_parameter_validation,
-                                mock_download_data, mock_create_genome_fasta,
+                                mock_download_data, mock_get_genome_sequences,
                                 mock_get_taxonomy):
-        # Mocked return values for the external functions
-        mock_parameter_validation.return_value = "mocked_rql_query"
-        mock_download_data.return_value = [
-            {'id': 'genome1', 'sequence': 'ATGC'},
-            {'id': 'genome2', 'sequence': 'GCTA'}]
-        mock_create_genome_fasta.return_value = MagicMock(
-            name='GenomeSequencesDirectoryFormat')
-        mock_get_taxonomy.return_value = MagicMock(name='TSVTaxonomyFormat')
 
         # Call the function
         get_bv_brc_genomes(
@@ -764,3 +756,29 @@ def test_get_sequences(self, mock_open, mock_download_data):
 
         # Check if the correct sequences were written to the proteins file
         mock_open().write.assert_any_call('>feature1\nMKV\n')
+
+
+class TestGetGenomeSequences(TestPluginBase):
+    package = 'rescript.tests'
+
+    @patch('rescript.bv_brc.download_data')
+    @patch('rescript.bv_brc.create_genome_fasta')
+    def test_get_genome_sequences(self, mock_create_genome_fasta,
+                                  mock_download_data):
+        # Sample response_genomes to be used as input
+        response_genomes = [
+            {'genome_id': '12345'},
+            {'genome_id': '67890'}
+        ]
+
+        # Call the function
+        get_genome_sequences(response_genomes)
+
+        # Assert that download_data was called with the correct arguments
+        mock_download_data.assert_called_once_with(
+            data_type="genome_sequence",
+            query=unittest.mock.ANY,
+            accept="application/json",
+            select=["accession", "description", "genome_name", "genome_id",
+                    "sequence"]
+        )