1 files changed, 48 insertions, 0 deletions
diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py
index 97365f4..51b4921 100644
--- a/tests/unit/test_llm.py
+++ b/tests/unit/test_llm.py
@@ -1,5 +1,6 @@
 """Test cases for procedures defined in llms """
 import pytest
+from gn3.llms.process import fetch_pubmed
 from gn3.llms.process import parse_context
 from gn3.llms.process import format_bibliography_info
 
@@ -54,3 +55,50 @@ def test_format_bib_info():
     assert all([format_bibliography_info(data) == expected
                 for data, expected
                 in zip(mock_fahamu_bib_info, expected_result)])
+
+
+@pytest.mark.unit_test
+def test_fetching_pubmed_info(monkeypatch):
+
+    def mock_load_file(filename, dir_path):
+        return {
+            "12121" : {
+                "Abstract": "items1",
+                "Author": "A1"
+            }
+        }
+    # patch the module with the mocked function
+
+    monkeypatch.setattr("gn3.llms.process.load_file", mock_load_file)
+    expected_results = [
+        {
+            "title": "Genes",
+            "year": "2014",
+            "doi": "https/article/genes/12121",
+            "doc_id": "12121",
+            "pubmed": {
+                "Abstract": "items1",
+                "Author": "A1"
+            }
+        },
+        {
+            "title": "Aging",
+            "year" : "2014",
+            "doc_id": "12122"
+        }
+    ]
+
+    data =  [        {
+            "title": "Genes",
+            "year": "2014",
+            "doi": "https/article/genes/12121",
+            "doc_id": "12121",
+        },
+        {
+            "title": "Aging",
+            "year" : "2014",
+            "doc_id": "12122"
+        }]
+
+    assert (fetch_pubmed(data, "/pubmed.json",  "data/")
+            == expected_results)