Added db_utils tests (#53)

* Added db_utils tests * Run new tests * Fixed merge issues * Review * Added new db util tests --------- Co-authored-by: kanesoban <[email protected]>
BlueBrain · Dec 17, 2024 · 29f8e8a · 29f8e8a
1 parent d5a2ccf
commit 29f8e8a
Show file tree

Hide file tree

Showing 8 changed files with 582 additions and 6 deletions.
diff --git a/.github/workflows/ci.yaml b/.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/
diff --git a/CHANGELOG.md b/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
+- Unit tests for database
 
 ### Fixed
 - Migrate LLM Evaluation logic to scripts and add tests

diff --git a/swarm_copy/tools/bluenaas_memodel_getall.py b/swarm_copy/tools/bluenaas_memodel_getall.py
@@ -29,7 +29,7 @@ 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(
+    memodel_type: Literal["single-neuron-simulation", "synaptome-simulation"] = Field(
         default="single-neuron-simulation",
         description="Type of simulation to retrieve.",
     )
@@ -55,7 +55,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.memodel_type,
                 "offset": self.input_schema.offset,
                 "page_size": self.input_schema.page_size,
             },

diff --git a/swarm_copy/tools/bluenaas_memodel_getone.py b/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(
+    memodel_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.memodel_id)}",
             headers={"Authorization": f"Bearer {self.metadata.token}"},
         )
 

diff --git a/swarm_copy/tools/traces_tool.py b/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):

diff --git a/swarm_copy_tests/app/database/__init__.py b/swarm_copy_tests/app/database/__init__.py
@@ -0,0 +1 @@
+"""Unit tests for database."""