uniprot litle bug fix

ur-whitelab · Jun 29, 2024 · 73d6fe6 · 73d6fe6
1 parent fb9b382
commit 73d6fe6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/mdagent/tools/base_tools/preprocess_tools/uniprot.py b/mdagent/tools/base_tools/preprocess_tools/uniprot.py
@@ -697,7 +697,7 @@ def get_ids(
             entry["primaryAccession"] for entry in accession
         ] if accession else []
         if single_id:
-            return all_ids.pop()
+            return [all_ids[0]] if all_ids else []
         return list(set(all_ids))
 
     def get_gene_names(self, query: str, primary_accession: str | None = None) -> list:

diff --git a/tests/test_preprocess/test_uniprot.py b/tests/test_preprocess/test_uniprot.py
@@ -478,7 +478,7 @@ def test_get_structure_info(query_uniprot):
     )
 
 
-def get_ids(query_uniprot):
+def test_get_ids(query_uniprot):
     hg_ids = [
         "P84792",
         "P02042",
@@ -508,7 +508,8 @@ def get_ids(query_uniprot):
     ]
     all_ids = query_uniprot.get_ids("hemoglobin")
     single_id = query_uniprot.get_ids("hemoglobin", single_id=True)
-    assert single_id in hg_ids
+    assert single_id[0] in hg_ids
+    assert len(single_id) == 1
     assert all(i in all_ids for i in hg_ids)