1 files changed, 57 insertions, 0 deletions
diff --git a/benchmark/tests/test_get_roots.py b/benchmark/tests/test_get_roots.py
new file mode 100644
index 000000000..dcc4327c5
--- /dev/null
+++ b/benchmark/tests/test_get_roots.py
@@ -0,0 +1,57 @@
+from agbenchmark.utils.dependencies.graphs import get_roots
+
+
+def test_get_roots():
+    graph = {
+        "nodes": [
+            {"id": "A", "data": {"category": []}},
+            {"id": "B", "data": {"category": []}},
+            {"id": "C", "data": {"category": []}},
+            {"id": "D", "data": {"category": []}},
+        ],
+        "edges": [
+            {"from": "A", "to": "B"},
+            {"from": "B", "to": "C"},
+        ],
+    }
+
+    result = get_roots(graph)
+    assert set(result) == {
+        "A",
+        "D",
+    }, f"Expected roots to be 'A' and 'D', but got {result}"
+
+
+def test_no_roots():
+    fully_connected_graph = {
+        "nodes": [
+            {"id": "A", "data": {"category": []}},
+            {"id": "B", "data": {"category": []}},
+            {"id": "C", "data": {"category": []}},
+        ],
+        "edges": [
+            {"from": "A", "to": "B"},
+            {"from": "B", "to": "C"},
+            {"from": "C", "to": "A"},
+        ],
+    }
+
+    result = get_roots(fully_connected_graph)
+    assert not result, "Expected no roots, but found some"
+
+
+# def test_no_rcoots():
+#     fully_connected_graph = {
+#         "nodes": [
+#             {"id": "A", "data": {"category": []}},
+#             {"id": "B", "data": {"category": []}},
+#             {"id": "C", "data": {"category": []}},
+#         ],
+#         "edges": [
+#             {"from": "A", "to": "B"},
+#             {"from": "D", "to": "C"},
+#         ],
+#     }
+#
+#     result = get_roots(fully_connected_graph)
+#     assert set(result) == {"A"}, f"Expected roots to be 'A', but got {result}"