Add more tests for aragorn functions

codecov.yml

@@ -1,3 +1,9 @@
 comment:
   layout: "header, files, footer"
-  hide_project_coverage: false
+  hide_project_coverage: false
+
+coverage:
+  status:
+    patch:
+      default:
+        enabled: false

tests/unit/aragorn/test_aragorn.py

@@ -0,0 +1,41 @@
+import copy
+import json
+import logging
+
+import pytest
+
+from shepherd_utils.broker import get_task
+from tests.helpers.generate_messages import creative_query
+from workers.aragorn.worker import aragorn
+
+
+@pytest.mark.asyncio
+async def test_aragorn_entrypoint(redis_mock, mocker):
+    """Test that Aragorn Lookup sends and gets back all the queries."""
+    mock_callback_response = mocker.patch("workers.aragorn.worker.get_message")
+    mock_callback_response.return_value = copy.deepcopy(creative_query)
+    logger = logging.getLogger(__name__)
+
+    await aragorn(
+        [
+            "test",
+            {
+                "query_id": "test",
+            },
+        ],
+        logger,
+    )
+
+    # Get the task that the ara should have put on the queue
+    task = await get_task("aragorn.lookup", "consumer", "test", logger)
+    assert task is not None
+    workflow = json.loads(task[1]["workflow"])
+    # make sure the workflow was correctly passed
+    assert len(workflow) == 5
+    assert [
+        "aragorn.lookup",
+        "aragorn.score",
+        "sort_results_score",
+        "filter_results_top_n",
+        "filter_kgraph_orphans",
+    ] == [op["id"] for op in workflow]

tests/unit/aragorn/test_aragorn_lookup.py

@@ -0,0 +1,48 @@
+import copy
+import json
+import logging
+
+import pytest
+
+from shepherd_utils.broker import get_task
+from tests.helpers.generate_messages import creative_query
+from workers.aragorn_lookup.worker import aragorn_lookup
+
+
+@pytest.mark.asyncio
+async def test_aragorn_creative_lookup(redis_mock, mocker):
+    """Test that Aragorn Lookup sends and gets back all the queries."""
+    mock_callback_response = mocker.patch("workers.aragorn_lookup.worker.get_message")
+    mock_callback_response.return_value = copy.deepcopy(creative_query)
+    mocker.patch("workers.aragorn_lookup.worker.save_message")
+    mocker.patch("workers.aragorn_lookup.worker.add_callback_id")
+    mock_running_callbacks = mocker.patch(
+        "workers.aragorn_lookup.worker.get_running_callbacks"
+    )
+    mock_running_callbacks.return_value = []
+    mock_response = mocker.Mock()
+    mocker.patch("httpx.AsyncClient.post", return_value=mock_response)
+    logger = logging.getLogger(__name__)
+
+    await aragorn_lookup(
+        [
+            "test",
+            {
+                "query_id": "test",
+                "workflow": json.dumps(
+                    [{"id": "aragorn.lookup"}, {"id": "aragorn.score"}]
+                ),
+            },
+        ],
+        logger,
+    )
+
+    # Get the task that the ara should have put on the queue
+    task = await get_task("aragorn.score", "consumer", "test", logger)
+    assert task is not None
+    workflow = json.loads(task[1]["workflow"])
+    # make sure the workflow was correctly passed
+    assert len(workflow) == 1
+    assert [
+        "aragorn.score",
+    ] == [op["id"] for op in workflow]

workers/aragorn_lookup/worker.py

@@ -85,8 +85,13 @@ async def aragorn_lookup(task, logger: logging.Logger):
         await save_message(
             f"{callback_id}_query_graph", message["message"]["query_graph"], logger
         )
-        # TODO: send single query to retriever
-        pass
+        message["callback"] = f"http://localhost:5439/callback/{callback_id}"
+
+        async with httpx.AsyncClient(timeout=100) as client:
+            await client.post(
+                "http://host.docker.internal:5781/asyncquery",
+                json=message,
+            )
     else:
         expanded_messages = expand_aragorn_query(message)
         # send all messages to retriever