Added app_utils tests

.github/workflows/ci.yaml

@@ -90,4 +90,4 @@ jobs:
           mypy src/ swarm_copy/
           # Include src/ directory in Python path to prioritize local files in pytest
           export PYTHONPATH=$(pwd)/src:$PYTHONPATH
-          pytest --color=yes
+          pytest --color=yes tests/ swarm_copy_tests/

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Tool implementations without langchain or langgraph dependencies
 - CRUDs.
 - BlueNaas CRUD tools
+- app unit tests
 
 ### Fixed
 - Migrate LLM Evaluation logic to scripts and add tests

swarm_copy/tools/bluenaas_memodel_getall.py

@@ -29,9 +29,11 @@ class InputMEModelGetAll(BaseModel):
     page_size: int = Field(
         default=20, description="Number of results returned by the API."
     )
-    model_type: Literal["single-neuron-simulation", "synaptome-simulation"] = Field(
-        default="single-neuron-simulation",
-        description="Type of simulation to retrieve.",
+    simulation_type: Literal["single-neuron-simulation", "synaptome-simulation"] = (
+        Field(
+            default="single-neuron-simulation",
+            description="Type of simulation to retrieve.",
+        )
     )
 
 
@@ -55,7 +57,7 @@ async def arun(self) -> PaginatedResponseUnionMEModelResponseSynaptomeModelRespo
         response = await self.metadata.httpx_client.get(
             url=f"{self.metadata.bluenaas_url}/neuron-model/{self.metadata.vlab_id}/{self.metadata.project_id}/me-models",
             params={
-                "simulation_type": self.input_schema.model_type,
+                "simulation_type": self.input_schema.simulation_type,
                 "offset": self.input_schema.offset,
                 "page_size": self.input_schema.page_size,
             },

swarm_copy/tools/bluenaas_memodel_getone.py

@@ -24,7 +24,7 @@ class MEModelGetOneMetadata(BaseMetadata):
 class InputMEModelGetOne(BaseModel):
     """Inputs for the BlueNaaS single-neuron simulation."""
 
-    model_id: str = Field(
+    simulation_id: str = Field(
         description="ID of the model to retrieve. Should be an https link."
     )
 
@@ -45,7 +45,7 @@ async def arun(self) -> MEModelResponse:
         )
 
         response = await self.metadata.httpx_client.get(
-            url=f"{self.metadata.bluenaas_url}/neuron-model/{self.metadata.vlab_id}/{self.metadata.project_id}/{quote_plus(self.input_schema.model_id)}",
+            url=f"{self.metadata.bluenaas_url}/neuron-model/{self.metadata.vlab_id}/{self.metadata.project_id}/{quote_plus(self.input_schema.simulation_id)}",
             headers={"Authorization": f"Bearer {self.metadata.token}"},
         )
 

swarm_copy/tools/traces_tool.py

@@ -1,6 +1,7 @@
 """Traces tool."""
 
 import logging
+from pathlib import Path
 from typing import Any, ClassVar
 
 from pydantic import BaseModel, Field
@@ -46,7 +47,7 @@ class GetTracesMetadata(BaseMetadata):
     knowledge_graph_url: str
     token: str
     trace_search_size: int
-    brainregion_path: str
+    brainregion_path: str | Path
 
 
 class GetTracesTool(BaseTool):

swarm_copy_tests/app/test_app_utils.py

@@ -0,0 +1,75 @@
+"""Test app utils."""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from fastapi.exceptions import HTTPException
+from httpx import AsyncClient
+
+from swarm_copy.app.app_utils import setup_engine, validate_project
+from swarm_copy.app.config import Settings
+
+
+@pytest.mark.asyncio
+async def test_validate_project(patch_required_env, httpx_mock, monkeypatch):
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "true")
+    httpx_client = AsyncClient()
+    token = "fake_token"
+    test_vp = {"vlab_id": "test_vlab_DB", "project_id": "project_id_DB"}
+    vlab_url = "https://openbluebrain.com/api/virtual-lab-manager/virtual-labs"
+
+    # test with bad config
+    httpx_mock.add_response(
+        url=f'{vlab_url}/{test_vp["vlab_id"]}/projects/{test_vp["project_id"]}',
+        status_code=404,
+    )
+    with pytest.raises(HTTPException) as error:
+        await validate_project(
+            httpx_client=httpx_client,
+            vlab_id=test_vp["vlab_id"],
+            project_id=test_vp["project_id"],
+            token=token,
+            vlab_project_url=vlab_url,
+        )
+    assert error.value.status_code == 401
+
+    # test with good config
+    httpx_mock.add_response(
+        url=f'{vlab_url}/{test_vp["vlab_id"]}/projects/{test_vp["project_id"]}',
+        json="test_project_ID",
+    )
+    await validate_project(
+        httpx_client=httpx_client,
+        vlab_id=test_vp["vlab_id"],
+        project_id=test_vp["project_id"],
+        token=token,
+        vlab_project_url=vlab_url,
+    )
+    # we jsut want to assert that the httpx_mock was called.
+
+
+@patch("neuroagent.app.app_utils.create_async_engine")
+def test_setup_engine(create_engine_mock, monkeypatch, patch_required_env):
+    create_engine_mock.return_value = AsyncMock()
+
+    monkeypatch.setenv("NEUROAGENT_DB__PREFIX", "prefix")
+
+    settings = Settings()
+
+    connection_string = "postgresql+asyncpg://user:password@localhost/dbname"
+    retval = setup_engine(settings=settings, connection_string=connection_string)
+    assert retval is not None
+
+
+@patch("neuroagent.app.app_utils.create_async_engine")
+def test_setup_engine_no_connection_string(
+    create_engine_mock, monkeypatch, patch_required_env
+):
+    create_engine_mock.return_value = AsyncMock()
+
+    monkeypatch.setenv("NEUROAGENT_DB__PREFIX", "prefix")
+
+    settings = Settings()
+
+    retval = setup_engine(settings=settings, connection_string=None)
+    assert retval is None

swarm_copy_tests/app/test_config.py

@@ -0,0 +1,71 @@
+"""Test config"""
+
+import pytest
+from pydantic import ValidationError
+
+from swarm_copy.app.config import Settings
+
+
+def test_required(monkeypatch, patch_required_env):
+    settings = Settings()
+
+    assert settings.tools.literature.url == "https://fake_url"
+    assert settings.knowledge_graph.base_url == "https://fake_url/api/nexus/v1"
+    assert settings.openai.token.get_secret_value() == "dummy"
+
+    # make sure not case sensitive
+    monkeypatch.delenv("NEUROAGENT_TOOLS__LITERATURE__URL")
+    monkeypatch.setenv("neuroagent_tools__literature__URL", "https://new_fake_url")
+
+    settings = Settings()
+    assert settings.tools.literature.url == "https://new_fake_url"
+
+
+def test_no_settings():
+    # We get an error when no custom variables provided
+    with pytest.raises(ValidationError):
+        Settings()
+
+
+def test_setup_tools(monkeypatch, patch_required_env):
+    monkeypatch.setenv("NEUROAGENT_TOOLS__TRACE__SEARCH_SIZE", "20")
+    monkeypatch.setenv("NEUROAGENT_TOOLS__MORPHO__SEARCH_SIZE", "20")
+    monkeypatch.setenv("NEUROAGENT_TOOLS__KG_MORPHO_FEATURES__SEARCH_SIZE", "20")
+
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__USERNAME", "user")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__PASSWORD", "pass")
+
+    settings = Settings()
+
+    assert settings.tools.morpho.search_size == 20
+    assert settings.tools.trace.search_size == 20
+    assert settings.tools.kg_morpho_features.search_size == 20
+    assert settings.keycloak.username == "user"
+    assert settings.keycloak.password.get_secret_value() == "pass"
+
+
+def test_check_consistency(monkeypatch):
+    # We get an error when no custom variables provided
+    url = "https://fake_url"
+    monkeypatch.setenv("NEUROAGENT_TOOLS__LITERATURE__URL", url)
+    monkeypatch.setenv("NEUROAGENT_KNOWLEDGE_GRAPH__URL", url)
+
+    with pytest.raises(ValueError):
+        Settings()
+
+    monkeypatch.setenv("NEUROAGENT_GENERATIVE__OPENAI__TOKEN", "dummy")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "true")
+
+    with pytest.raises(ValueError):
+        Settings()
+
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "false")
+
+    with pytest.raises(ValueError):
+        Settings()
+
+    monkeypatch.setenv("NEUROAGENT_KNOWLEDGE_GRAPH__BASE_URL", "http://fake_nexus.com")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "true")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__PASSWORD", "Hello")
+
+    Settings()

swarm_copy_tests/app/test_dependencies.py

@@ -0,0 +1,171 @@
+"""Test dependencies."""
+
+import json
+import os
+from pathlib import Path
+from typing import AsyncIterator
+from unittest.mock import Mock
+
+import pytest
+from httpx import AsyncClient
+
+from swarm_copy.app.dependencies import (
+    Settings,
+    get_cell_types_kg_hierarchy,
+    get_connection_string,
+    get_httpx_client,
+    get_settings,
+    get_update_kg_hierarchy,
+    get_user_id,
+)
+
+
+def test_get_settings(patch_required_env):
+    settings = get_settings()
+    assert settings.tools.literature.url == "https://fake_url"
+    assert settings.knowledge_graph.url == "https://fake_url/api/nexus/v1/search/query/"
+
+
+@pytest.mark.asyncio
+async def test_get_httpx_client():
+    request = Mock()
+    request.headers = {"x-request-id": "greatid"}
+    httpx_client_iterator = get_httpx_client(request=request)
+    assert isinstance(httpx_client_iterator, AsyncIterator)
+    async for httpx_client in httpx_client_iterator:
+        assert isinstance(httpx_client, AsyncClient)
+        assert httpx_client.headers["x-request-id"] == "greatid"
+
+
+@pytest.mark.asyncio
+async def test_get_user(httpx_mock, monkeypatch, patch_required_env):
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__USERNAME", "fake_username")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__PASSWORD", "fake_password")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__ISSUER", "https://great_issuer.com")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "true")
+
+    fake_response = {
+        "sub": "12345",
+        "email_verified": False,
+        "name": "Machine Learning Test User",
+        "groups": [],
+        "preferred_username": "sbo-ml",
+        "given_name": "Machine Learning",
+        "family_name": "Test User",
+        "email": "[email protected]",
+    }
+    httpx_mock.add_response(
+        url="https://great_issuer.com/protocol/openid-connect/userinfo",
+        json=fake_response,
+    )
+
+    settings = Settings()
+    client = AsyncClient()
+    token = "eyJgreattoken"
+    user_id = await get_user_id(token=token, settings=settings, httpx_client=client)
+
+    assert user_id == fake_response["sub"]
+
+
+@pytest.mark.asyncio
+async def test_get_update_kg_hierarchy(
+    tmp_path, httpx_mock, monkeypatch, patch_required_env
+):
+    token = "fake_token"
+    file_name = "fake_file"
+    client = AsyncClient()
+
+    file_url = "https://fake_file_url"
+
+    monkeypatch.setenv(
+        "NEUROAGENT_KNOWLEDGE_GRAPH__HIERARCHY_URL", "http://fake_hierarchy_url.com"
+    )
+
+    settings = Settings(
+        knowledge_graph={"br_saving_path": tmp_path / "test_brain_region.json"}
+    )
+
+    json_response_url = {
+        "head": {"vars": ["file_url"]},
+        "results": {"bindings": [{"file_url": {"type": "uri", "value": file_url}}]},
+    }
+    with open(
+        Path(__file__).parent.parent.parent
+        / "tests"
+        / "data"
+        / "KG_brain_regions_hierarchy_test.json"
+    ) as fh:
+        json_response_file = json.load(fh)
+
+    httpx_mock.add_response(
+        url=settings.knowledge_graph.sparql_url, json=json_response_url
+    )
+    httpx_mock.add_response(url=file_url, json=json_response_file)
+
+    await get_update_kg_hierarchy(
+        token,
+        client,
+        settings,
+        file_name,
+    )
+
+    assert os.path.exists(settings.knowledge_graph.br_saving_path)
+
+
+@pytest.mark.asyncio
+async def test_get_cell_types_kg_hierarchy(
+    tmp_path, httpx_mock, monkeypatch, patch_required_env
+):
+    token = "fake_token"
+    file_name = "fake_file"
+    client = AsyncClient()
+
+    file_url = "https://fake_file_url"
+    monkeypatch.setenv(
+        "NEUROAGENT_KNOWLEDGE_GRAPH__HIERARCHY_URL", "http://fake_hierarchy_url.com"
+    )
+
+    settings = Settings(
+        knowledge_graph={"ct_saving_path": tmp_path / "test_cell_types_region.json"}
+    )
+
+    json_response_url = {
+        "head": {"vars": ["file_url"]},
+        "results": {"bindings": [{"file_url": {"type": "uri", "value": file_url}}]},
+    }
+    with open(
+        Path(__file__).parent.parent.parent
+        / "tests"
+        / "data"
+        / "kg_cell_types_hierarchy_test.json"
+    ) as fh:
+        json_response_file = json.load(fh)
+
+    httpx_mock.add_response(
+        url=settings.knowledge_graph.sparql_url, json=json_response_url
+    )
+    httpx_mock.add_response(url=file_url, json=json_response_file)
+
+    await get_cell_types_kg_hierarchy(
+        token,
+        client,
+        settings,
+        file_name,
+    )
+
+    assert os.path.exists(settings.knowledge_graph.ct_saving_path)
+
+
+def test_get_connection_string_full(monkeypatch, patch_required_env):
+    monkeypatch.setenv("NEUROAGENT_DB__PREFIX", "http://")
+    monkeypatch.setenv("NEUROAGENT_DB__USER", "John")
+    monkeypatch.setenv("NEUROAGENT_DB__PASSWORD", "Doe")
+    monkeypatch.setenv("NEUROAGENT_DB__HOST", "localhost")
+    monkeypatch.setenv("NEUROAGENT_DB__PORT", "5000")
+    monkeypatch.setenv("NEUROAGENT_DB__NAME", "test")
+
+    settings = Settings()
+    result = get_connection_string(settings)
+    assert (
+        result == "http://John:Doe@localhost:5000/test"
+    ), "must return fully formed connection string"