Fixed some failing tests

BlueBrain · Dec 18, 2024 · f81ab8d · f81ab8d
1 parent e1a9095
commit f81ab8d
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 28 deletions.
diff --git a/src/neuroagent/cell_types.py b/src/neuroagent/cell_types.py
@@ -17,7 +17,6 @@ class CellTypesMeta:
     """
 
     def __init__(self) -> None:
-        self.name_: dict[str, str | None] = {}
         self.name_: dict[Any, Any | None] = {}
         self.descendants_ids: dict[str, set[str]] = {}
 

diff --git a/swarm_copy_tests/tools/test_electrophys_tool.py b/swarm_copy_tests/tools/test_electrophys_tool.py
@@ -60,11 +60,10 @@ async def test_arun(self, httpx_mock):
                 token="fake_token")
         )
         response = await tool.arun()
-        assert isinstance(response, FeatureOutput)
-        assert isinstance(response.feature_dict, dict)
-        assert len(response.feature_dict.keys()) == 1
+        assert isinstance(response, dict)
+        assert len(response.keys()) == 1
         assert (
-            len(response.feature_dict["step_0"].keys())
+            len(response["step_0"].keys())
             == 2  # mean_frequency + 1 for stimulus current added manually
         )
 
@@ -112,11 +111,10 @@ async def test_arun_with_amplitude(self, httpx_mock):
                 token="fake_token")
         )
         response = await tool.arun()
-        assert isinstance(response, FeatureOutput)
-        assert isinstance(response.feature_dict, dict)
-        assert len(response.feature_dict.keys()) == 1
+        assert isinstance(response, dict)
+        assert len(response.keys()) == 1
         assert (
-            len(response.feature_dict["step_0.25"].keys())
+            len(response["step_0.25"].keys())
             == 2  # mean_frequency + 1 for stimulus current added manually
         )
 
@@ -164,11 +162,10 @@ async def test_arun_without_stimuli_types(self, httpx_mock):
 
         # Without stimuli types and calculated features
         response = await tool.arun()
-        assert isinstance(response, FeatureOutput)
-        assert isinstance(response.feature_dict, dict)
-        assert len(response.feature_dict.keys()) == 1
+        assert isinstance(response, dict)
+        assert len(response.keys()) == 1
         assert (
-            len(response.feature_dict["step_0"].keys())
+            len(response["step_0"].keys())
             == len(list(CALCULATED_FEATURES.__args__[0].__args__))
             + 1  # 1 for stimulus current added manually
         )

diff --git a/swarm_copy_tests/tools/test_get_morpho_tool.py b/swarm_copy_tests/tools/test_get_morpho_tool.py
@@ -41,7 +41,7 @@ async def test_arun(self, httpx_mock, brain_region_json_path, tmp_path):
         response = await tool.arun()
         assert isinstance(response, list)
         assert len(response) == 2
-        assert isinstance(response[0], KnowledgeGraphOutput)
+        assert isinstance(response[0], dict)
 
     @pytest.mark.asyncio
     async def test_arun_errors(self, httpx_mock, brain_region_json_path, tmp_path):

diff --git a/swarm_copy_tests/tools/test_kg_morpho_features_tool.py b/swarm_copy_tests/tools/test_kg_morpho_features_tool.py
@@ -51,7 +51,7 @@ async def test_arun(self, httpx_mock, brain_region_json_path):
         response = await tool.arun()
         assert isinstance(response, list)
         assert len(response) == 2
-        assert isinstance(response[0], KGMorphoFeatureOutput)
+        assert isinstance(response[0], dict)
 
     @pytest.mark.asyncio
     async def test_arun_errors(self, httpx_mock, brain_region_json_path):

diff --git a/swarm_copy_tests/tools/test_literature_search_tool.py b/swarm_copy_tests/tools/test_literature_search_tool.py
@@ -46,4 +46,4 @@ async def test_arun(self, httpx_mock):
         response = await tool.arun()
         assert isinstance(response, list)
         assert len(response) == reranker_k
-        assert isinstance(response[0], ParagraphMetadata)
+        assert isinstance(response[0], dict)
diff --git a/swarm_copy_tests/tools/test_morphology_features_tool.py b/swarm_copy_tests/tools/test_morphology_features_tool.py
@@ -49,8 +49,8 @@ async def test_arun(self, httpx_mock):
         )
 
         response = await tool.arun()
-        assert isinstance(response[0], MorphologyFeatureOutput)
-        assert len(response[0].feature_dict) == 23
+        assert isinstance(response[0], dict)
+        assert len(response[0]) == 23
 
     @pytest.mark.asyncio
     async def test_arun_errors_404(self, httpx_mock):

diff --git a/swarm_copy_tests/tools/test_resolve_entities_tool.py b/swarm_copy_tests/tools/test_resolve_entities_tool.py
@@ -74,27 +74,27 @@ async def test_arun(httpx_mock, get_resolve_query_output):
         BRResolveOutput(
             brain_region_name="Field CA1",
             brain_region_id="http://api.brain-map.org/api/v2/data/Structure/382",
-        ),
+        ).model_dump(),
         BRResolveOutput(
             brain_region_name="Field CA2",
             brain_region_id="http://api.brain-map.org/api/v2/data/Structure/423",
-        ),
+        ).model_dump(),
         BRResolveOutput(
             brain_region_name="Field CA3",
             brain_region_id="http://api.brain-map.org/api/v2/data/Structure/463",
-        ),
+        ).model_dump(),
         MTypeResolveOutput(
             mtype_name="Interneuron", mtype_id="https://neuroshapes.org/Interneuron"
-        ),
+        ).model_dump(),
         MTypeResolveOutput(
             mtype_name="Hippocampus CA3 Oriens Interneuron",
             mtype_id="http://uri.interlex.org/base/ilx_0105044",
-        ),
+        ).model_dump(),
         MTypeResolveOutput(
             mtype_name="Spinal Cord Ventral Horn Interneuron IA",
             mtype_id="http://uri.interlex.org/base/ilx_0110929",
-        ),
+        ).model_dump(),
         EtypeResolveOutput(
             etype_name="bAC", etype_id="http://uri.interlex.org/base/ilx_0738199"
-        ),
+        ).model_dump(),
     ]
diff --git a/swarm_copy_tests/tools/test_traces_tool.py b/swarm_copy_tests/tools/test_traces_tool.py
@@ -40,8 +40,8 @@ async def test_arun(self, httpx_mock, brain_region_json_path):
         response = await tool.arun()
         assert isinstance(response, list)
         assert len(response) == 2
-        assert isinstance(response[0], TracesOutput)
-        assert isinstance(response[0], TracesOutput)
+        assert isinstance(response[0], dict)
+        assert isinstance(response[0], dict)
 
     @pytest.mark.httpx_mock(can_send_already_matched_responses=True)
     @pytest.mark.asyncio
@@ -71,7 +71,7 @@ async def test_arun_with_etype(self, httpx_mock, brain_region_json_path):
         response = await tool.arun()
         assert isinstance(response, list)
         assert len(response) == 2
-        assert isinstance(response[0], TracesOutput)
+        assert isinstance(response[0], dict)
 
     @pytest.mark.asyncio
     async def test_arun_errors(self, httpx_mock, brain_region_json_path):