ci: add tests for distribute work

.gitlab-ci.yml

@@ -40,6 +40,8 @@ build-coyote:
 test-coyote:
   stage: test
   image: "${CI_REGISTRY_IMAGE}:${CI_COMMIT_REF_SLUG}"
+  variables:
+    PYTHONPATH: "${CI_PROJECT_DIR}/build/python:${CI_PROJECT_DIR}/apps"
   tags:
     - docker-any-image, dkrz
   script:

apps/hiopy/CMakeLists.txt

@@ -3,3 +3,5 @@ install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}"
   EXCLUDE_FROM_ALL
   DESTINATION "."
 )
+
+add_subdirectory(tests)

apps/hiopy/tests/CMakeLists.txt

+add_test(
+  NAME hiopy.test_distribute_work
+  COMMAND "${CMAKE_CURRENT_SOURCE_DIR}/test_distribute_work.py"
+)

apps/hiopy/tests/test_distribute_work.py

+#!/usr/bin/env python3
+
+import numpy as np
+from hiopy._distribute_work import distribute_work
+try:
+    from pprint import pp
+except (ImportError):
+    print("Cannot find pprint falling back to print")
+    pp = print
+
+
+class DummyDataVar:
+    def __init__(self, shape, chunksize=None):
+        self.shape = shape
+        self.chunks = (1, chunksize or shape[1])
+        self.attrs = {"_ARRAY_DIMENSIONS": ["time", "cell"]}
+        self.dtype = np.dtype("float64")
+
+
+def test_sequential():
+    print("Test sequential case")
+    var1 = DummyDataVar([10, 10])
+    var2 = DummyDataVar([10, 10])
+    result = distribute_work({"g1": [var1],
+                              "g2": [var2]}, 1)
+    pp(result)
+    assert len(result) == 1
+    assert result[0] == [("g1", 0, slice(None), [var1]),
+                         ("g2", 0, slice(None), [var2])]
+
+
+def test_dist_groups():
+    print("Test distributing groups")
+    var1 = DummyDataVar([10, 10])
+    var2 = DummyDataVar([10, 10])
+    result = distribute_work({"g1": [var1],
+                              "g2": [var2]}, 2)
+    pp(result)
+    assert len(result) == 2
+    assert result[0] == [("g1", 0, slice(None), [var1])]
+    assert result[1] == [("g2", 0, slice(None), [var2])]
+
+
+def test_dist_vars():
+    print("Test distributing variables")
+    var1 = DummyDataVar([10, 10])
+    var2 = DummyDataVar([10, 10])
+    result = distribute_work({"g": [var1, var2]}, 2)
+    pp(result)
+    assert len(result) == 2
+    assert result[0] == [("g", 0, slice(None), [var1])]
+    assert result[1] == [("g", 1, slice(None), [var2])]
+
+
+def test_dist_chunks():
+    print("Test distributing chunks")
+    var = DummyDataVar([10, 10], 5)
+    result = distribute_work({"g": [var]}, 2)
+    pp(result)
+    assert len(result) == 2
+    assert result[0] == [("g", 0, slice(0, 5), [var])]
+    assert result[1] == [("g", 0, slice(5, 10), [var])]
+
+
+if __name__ == "__main__":
+    test_sequential()
+    test_dist_groups()
+    test_dist_vars()
+    test_dist_chunks()