From c037a2e003bee9be12b76947a5b8c1fa9e6312ec Mon Sep 17 00:00:00 2001
From: Weina Ji <weina.ji@epfl.ch>
Date: Fri, 9 Feb 2024 17:04:17 +0100
Subject: [PATCH] fix tests

---
 neurodamus/io/sonata_config.py               | 8 ++++----
 tests/integration-e2e/test_synapse_reader.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/neurodamus/io/sonata_config.py b/neurodamus/io/sonata_config.py
index d17043bcb..f01108e64 100644
--- a/neurodamus/io/sonata_config.py
+++ b/neurodamus/io/sonata_config.py
@@ -5,10 +5,10 @@
 import libsonata
 import logging
 import os.path
-from enum import Enum
+from enum import StrEnum
 
 
-class ConnectionTypes(Enum):
+class ConnectionTypes(StrEnum):
     Synaptic = "Synaptic"
     GapJunction = "GapJunction"
     NeuroModulation = "NeuroModulation"
@@ -332,10 +332,10 @@ def parsedProjections(self):
                     Type=projection_type_convert.get(pop_type)
                 )
                 # Reverse projection direction for Astrocyte projection: from neurons to astrocytes
-                if projection.get("Type") == "NeuroGlial":
+                if projection.get("Type") == ConnectionTypes.NeuroGlial:
                     projection["Source"], projection["Destination"] = projection["Destination"], \
                         projection["Source"]
-                if projection.get("Type") == "GlioVascular":
+                if projection.get("Type") == ConnectionTypes.GlioVascular:
                     for node_file_info in self._circuit_networks["nodes"]:
                         for _, pop_info in node_file_info["populations"].items():
                             if pop_info.get("type") == "vasculature":
diff --git a/tests/integration-e2e/test_synapse_reader.py b/tests/integration-e2e/test_synapse_reader.py
index c4cf07d24..72b927d49 100644
--- a/tests/integration-e2e/test_synapse_reader.py
+++ b/tests/integration-e2e/test_synapse_reader.py
@@ -8,7 +8,7 @@
 
 def test_gapjunction_sonata_reader():
     sonata_file = str(SIM_DIR / "mini_thalamus_sonata/gapjunction/edges.h5")
-    sonata_reader = GapJunctionSynapseReader.create(sonata_file, 1)
+    sonata_reader = GapJunctionSynapseReader.create(sonata_file)
     syn_params_sonata = sonata_reader._load_synapse_parameters(1)
     ref_junction_id_pre = np.array([10257., 43930., 226003., 298841., 324744.,
                                     1094745., 1167632., 1172523., 1260104.])