From 524c337f714aa0aa942259ceced07f3d0de4dc75 Mon Sep 17 00:00:00 2001
From: Matthieu Dorier <mdorier@anl.gov>
Date: Tue, 30 Jan 2024 15:54:01 +0000
Subject: [PATCH] added tests for ssg manager in python

---
 python/mochi/bedrock/server.py             | 16 ++++++++++------
 python/mochi/bedrock/spec.py               |  2 +-
 python/mochi/bedrock/test_abtio_manager.py |  2 +-
 python/mochi/bedrock/test_margo_manager.py |  2 +-
 src/SSGManager.cpp                         | 13 ++++++++++---
 5 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/python/mochi/bedrock/server.py b/python/mochi/bedrock/server.py
index a7181a4..0f8b123 100644
--- a/python/mochi/bedrock/server.py
+++ b/python/mochi/bedrock/server.py
@@ -160,7 +160,8 @@ def config(self):
 
     @property
     def spec(self) -> list[SSGSpec]:
-        return [SSGSpec.from_dict(group) for group in self._server.config]
+        abt_spec = self._server.margo.spec.argobots
+        return [SSGSpec.from_dict(group, abt_spec) for group in self.config]
 
     def resolve(self, margo, location: str) -> pymargo.core.Address:
         if isinstance(margo, pymargo.core.Engine):
@@ -177,12 +178,15 @@ def __len__(self) -> int:
     def __getitem__(self, key: int|str) -> SSGGroup:
         return SSGGroup(self._internal.get_group(key))
 
-    def create(self, name: str, config: str|dict|SSGSpec, pool: Pool,
-               bootstrap: str, group_file: str = "") -> SSGGroup:
-        if isinstance(config, dict):
-            config = json.dumps(config)
+    def create(self, name: str, pool: str|int|Pool = "__primary__",
+               config: str|dict|SSGSpec = "{}",
+               bootstrap: str = "init", group_file: str = "") -> SSGGroup:
+        if not isinstance(pool, Pool):
+            pool = self._server.margo.pools[pool]
+        if isinstance(config, str):
+            config = json.loads(config)
         elif isinstance(config, SSGSpec):
-            config = config.to_json()
+            config = config.to_dict()
         pool = pool._internal
         return SSGGroup(self._internal.create_group(name, config, pool, bootstrap, group_file))
 
diff --git a/python/mochi/bedrock/spec.py b/python/mochi/bedrock/spec.py
index 0c9910a..617ef42 100644
--- a/python/mochi/bedrock/spec.py
+++ b/python/mochi/bedrock/spec.py
@@ -1224,7 +1224,7 @@ def from_dict(data: dict, abt_spec: ArgobotsSpec) -> 'SSGSpec':
         """
         args = data.copy()
         args['pool'] = abt_spec.pools[data['pool']]
-        args['swim'] = SwimSpec.from_dict(**args['swim'])
+        args['swim'] = SwimSpec.from_dict(args['swim'])
         ssg = SSGSpec(**args)
         return ssg
 
diff --git a/python/mochi/bedrock/test_abtio_manager.py b/python/mochi/bedrock/test_abtio_manager.py
index 3a65a11..15d53a1 100644
--- a/python/mochi/bedrock/test_abtio_manager.py
+++ b/python/mochi/bedrock/test_abtio_manager.py
@@ -4,7 +4,7 @@
 import mochi.bedrock.spec as spec
 
 
-class TestMargoInstance(unittest.TestCase):
+class TestAbtIOManager(unittest.TestCase):
 
     def setUp(self):
         config = {
diff --git a/python/mochi/bedrock/test_margo_manager.py b/python/mochi/bedrock/test_margo_manager.py
index 9adf79d..2a44965 100644
--- a/python/mochi/bedrock/test_margo_manager.py
+++ b/python/mochi/bedrock/test_margo_manager.py
@@ -4,7 +4,7 @@
 import mochi.bedrock.spec as spec
 
 
-class TestMargoInstance(unittest.TestCase):
+class TestMargoManager(unittest.TestCase):
 
     def setUp(self):
         config = {
diff --git a/src/SSGManager.cpp b/src/SSGManager.cpp
index 885f35e..de4ee25 100644
--- a/src/SSGManager.cpp
+++ b/src/SSGManager.cpp
@@ -214,22 +214,29 @@ std::shared_ptr<NamedDependency> SSGManager::getGroup(const std::string& group_n
 #ifdef ENABLE_SSG
     auto it = std::find_if(self->m_ssg_groups.begin(), self->m_ssg_groups.end(),
                            [&](auto& g) { return g->getName() == group_name; });
-    if (it == self->m_ssg_groups.end())
+    if (it == self->m_ssg_groups.end()) {
+        throw DETAILED_EXCEPTION("Could not find SSG group with name \"{}\"", group_name);
         return nullptr;
-    else
+    } else {
         return *it;
+    }
 #else
     (void)group_name;
+    throw DETAILED_EXCEPTION("Bedrock was not compiler with SSG support");
     return nullptr;
 #endif
 }
 
 std::shared_ptr<NamedDependency> SSGManager::getGroup(uint32_t group_index) const {
 #ifdef ENABLE_SSG
-    if(group_index >= self->m_ssg_groups.size()) return nullptr;
+    if(group_index >= self->m_ssg_groups.size()) {
+        throw DETAILED_EXCEPTION("Could not find SSG group at index {}", group_index);
+        return nullptr;
+    }
     return self->m_ssg_groups[group_index];
 #else
     (void)group_index;
+    throw DETAILED_EXCEPTION("Bedrock was not compiler with SSG support");
     return nullptr;
 #endif
 }