From 97373bef3aaf44b0ad17c9342a1b752ed9025a96 Mon Sep 17 00:00:00 2001
From: Hugo Cachitas <hcachitas@gmail.com>
Date: Sun, 31 Mar 2024 00:25:42 +0000
Subject: [PATCH] Add module methods

---
 src/stringdb/__init__.py |  9 +++++++++
 tests/test_stringdb.py   | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 tests/test_stringdb.py

diff --git a/src/stringdb/__init__.py b/src/stringdb/__init__.py
index c23d9d2..3d4d376 100644
--- a/src/stringdb/__init__.py
+++ b/src/stringdb/__init__.py
@@ -12,3 +12,12 @@
 
 DEFAULT_CALLER_IDENTITY = f"{__name__} {__version__}"
 
+
+def map_identifiers(identifiers: List[str], species: int):
+    with Client() as client:
+        return client.map_identifiers(identifiers=identifiers, species=species)
+
+
+def interaction_partners(identifiers: List[str], species: int):
+    with Client() as client:
+        return client.interaction_partners(identifiers=identifiers, species=species)
diff --git a/tests/test_stringdb.py b/tests/test_stringdb.py
new file mode 100644
index 0000000..c48a92f
--- /dev/null
+++ b/tests/test_stringdb.py
@@ -0,0 +1,40 @@
+import stringdb
+
+
+def test_map_identifiers():
+    identifiers = stringdb.map_identifiers(["edin"], 7227)
+
+    assert len(identifiers) == 1
+
+    for id_ in identifiers:
+        keys = id_.keys()
+        assert "queryIndex" in keys
+        assert "queryItem" in keys
+        assert "stringId" in keys
+        assert "ncbiTaxonId" in keys
+        assert "taxonName" in keys
+        assert "preferredName" in keys
+        assert "annotation" in keys
+        assert "annotation" in keys
+
+
+def test_interaction_partners():
+    interaction_partners = stringdb.interaction_partners(["edin"], 7227)
+
+    assert len(interaction_partners) > 1
+
+    for id_ in interaction_partners:
+        keys = id_.keys()
+        assert "stringId_A" in keys
+        assert "stringId_B" in keys
+        assert "preferredName_A" in keys
+        assert "preferredName_B" in keys
+        assert "ncbiTaxonId" in keys
+        assert "score" in keys
+        assert "nscore" in keys
+        assert "fscore" in keys
+        assert "pscore" in keys
+        assert "ascore" in keys
+        assert "escore" in keys
+        assert "dscore" in keys
+        assert "tscore" in keys