From 552fae4b95f2410687e06d07e72c083720034b39 Mon Sep 17 00:00:00 2001 From: Artemy Belousov Date: Wed, 31 Jul 2024 10:14:30 +0300 Subject: [PATCH] Generalize common graph store tests (#68) --- motleycrew/common/__init__.py | 9 + motleycrew/common/enums.py | 14 ++ tests/test_storage/__init__.py | 21 ++ tests/test_storage/test_graph_store.py | 188 +++++++++++++++++ tests/test_storage/test_kuzu_graph_store.py | 212 +------------------- 5 files changed, 240 insertions(+), 204 deletions(-) create mode 100644 tests/test_storage/test_graph_store.py diff --git a/motleycrew/common/__init__.py b/motleycrew/common/__init__.py index 2543ae6a..3011bc95 100644 --- a/motleycrew/common/__init__.py +++ b/motleycrew/common/__init__.py @@ -18,4 +18,13 @@ "Defaults", "MotleySupportedTool", "MotleyAgentFactory", + "logger", + "configure_logging", + "AsyncBackend", + "GraphStoreType", + "LLMFamily", + "LLMFramework", + "LunaryEventName", + "LunaryRunType", + "TaskUnitStatus", ] diff --git a/motleycrew/common/enums.py b/motleycrew/common/enums.py index a095eca4..1fa438c4 100644 --- a/motleycrew/common/enums.py +++ b/motleycrew/common/enums.py @@ -5,21 +5,29 @@ class LLMFamily: OPENAI = "openai" ANTHROPIC = "anthropic" + ALL = {OPENAI, ANTHROPIC} + class LLMFramework: LANGCHAIN = "langchain" LLAMA_INDEX = "llama_index" + ALL = {LANGCHAIN, LLAMA_INDEX} + class GraphStoreType: KUZU = "kuzu" + ALL = {KUZU} + class TaskUnitStatus: PENDING = "pending" RUNNING = "running" DONE = "done" + ALL = {PENDING, RUNNING, DONE} + class LunaryRunType: LLM = "llm" @@ -28,6 +36,8 @@ class LunaryRunType: CHAIN = "chain" EMBED = "embed" + ALL = {LLM, AGENT, TOOL, CHAIN, EMBED} + class LunaryEventName: START = "start" @@ -35,6 +45,8 @@ class LunaryEventName: UPDATE = "update" ERROR = "error" + ALL = {START, END, UPDATE, ERROR} + class AsyncBackend: """Backends for parallel crew execution. @@ -48,3 +60,5 @@ class AsyncBackend: ASYNCIO = "asyncio" THREADING = "threading" NONE = "none" + + ALL = {ASYNCIO, THREADING, NONE} diff --git a/tests/test_storage/__init__.py b/tests/test_storage/__init__.py index e69de29b..684f66a3 100644 --- a/tests/test_storage/__init__.py +++ b/tests/test_storage/__init__.py @@ -0,0 +1,21 @@ +import kuzu +import pytest + +from motleycrew.storage import MotleyKuzuGraphStore + + +class GraphStoreFixtures: + @pytest.fixture + def kuzu_graph_store(self, tmpdir): + db_path = tmpdir / "test_db" + db = kuzu.Database(str(db_path)) + + graph_store = MotleyKuzuGraphStore(db) + return graph_store + + @pytest.fixture + def graph_store(self, request, kuzu_graph_store): + graph_stores = { + "kuzu": kuzu_graph_store, + } + return graph_stores.get(request.param) diff --git a/tests/test_storage/test_graph_store.py b/tests/test_storage/test_graph_store.py new file mode 100644 index 00000000..1af14613 --- /dev/null +++ b/tests/test_storage/test_graph_store.py @@ -0,0 +1,188 @@ +from typing import Optional + +import pytest + +from motleycrew.common import GraphStoreType +from motleycrew.storage import MotleyGraphNode +from motleycrew.storage import MotleyKuzuGraphStore +from tests.test_storage import GraphStoreFixtures + + +class Entity(MotleyGraphNode): + int_param: int + optional_str_param: Optional[str] = None + optional_list_str_param: Optional[list[str]] = None + + +class TestMotleyGraphStore(GraphStoreFixtures): + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_insert_new_node(self, graph_store): + entity = Entity(int_param=1) + created_entity = graph_store.insert_node(entity) + assert created_entity.id is not None + assert entity.id is not None # mutated in place + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_insert_node_and_retrieve(self, graph_store): + entity = Entity(int_param=1, optional_str_param="test", optional_list_str_param=["a", "b"]) + inserted_entity = graph_store.insert_node(entity) + assert inserted_entity.id is not None + + retrieved_entity = graph_store.get_node_by_class_and_id( + node_class=Entity, node_id=inserted_entity.id + ) + assert inserted_entity == retrieved_entity + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_check_node_exists_true(self, graph_store): + entity = Entity(int_param=1) + + graph_store.insert_node(entity) + assert graph_store.check_node_exists(entity) + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_check_node_exists_false(self, graph_store): + entity = Entity(int_param=1) + assert graph_store.check_node_exists(entity) is False + + MotleyKuzuGraphStore._set_node_id(node=entity, node_id=2) + assert graph_store.check_node_exists(entity) is False + + graph_store.ensure_node_table(Entity) + assert graph_store.check_node_exists(entity) is False + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_create_relation(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + graph_store.insert_node(entity1) + graph_store.insert_node(entity2) + + graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") + + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="p") + assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="q") + assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_upsert_triplet(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + graph_store.upsert_triplet(from_node=entity1, to_node=entity2, label="p") + + assert graph_store.check_node_exists(entity1) + assert graph_store.check_node_exists(entity2) + + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="p") + assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="q") + assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_nodes_do_not_exist(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + + assert not graph_store.check_node_exists(entity1) + assert not graph_store.check_node_exists(entity2) + + assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) + assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1, label="p") + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_relation_does_not_exist(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + + assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) + assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) + + graph_store.insert_node(entity1) + graph_store.insert_node(entity2) + + assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) + assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_delete_node(self, graph_store): + entity = Entity(int_param=1) + graph_store.insert_node(entity) + assert graph_store.check_node_exists(entity) is True + + graph_store.delete_node(entity) + assert graph_store.check_node_exists(entity) is False + + entity.int_param = 2 # check that entity is not frozen + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_delete_entity_with_relations(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + + graph_store.insert_node(entity1) + graph_store.insert_node(entity2) + graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) is True + + graph_store.delete_node(entity1) + assert graph_store.check_node_exists(entity1) is False + assert graph_store.check_node_exists(entity2) is True + assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) is False + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_set_property(self, graph_store): + entity = Entity(int_param=1) + graph_store.insert_node(entity) + assert entity.optional_str_param is None + assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param is None + + entity.optional_str_param = "test" + assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param == "test" + + entity.optional_list_str_param = ["a", "b"] + assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_list_str_param == [ + "a", + "b", + ] + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_run_cypher_query(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2) + + graph_store.insert_node(entity1) + graph_store.insert_node(entity2) + graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") + + query = """ + MATCH (a:Entity {int_param: 1})-[r]->(b:Entity {int_param: 2}) + RETURN a, r, b + """ + result = graph_store.run_cypher_query(query) + assert len(result) == 1 + assert len(result[0]) == 3 + + a, r, b = result[0] + assert a["int_param"] == 1 + assert b["int_param"] == 2 + assert r["_label"] == "p" + + @pytest.mark.parametrize("graph_store", GraphStoreType.ALL, indirect=True) + def test_run_cypher_query_with_container(self, graph_store): + entity1 = Entity(int_param=1) + entity2 = Entity(int_param=2, optional_list_str_param=["a", "b"]) + + graph_store.insert_node(entity1) + graph_store.insert_node(entity2) + + query = """ + MATCH (a:Entity) + WHERE a.int_param = 2 + RETURN a + """ + result = graph_store.run_cypher_query(query, container=Entity) + assert len(result) == 1 + assert isinstance(result[0], Entity) + assert result[0].int_param == 2 + assert result[0].optional_list_str_param == ["a", "b"] diff --git a/tests/test_storage/test_kuzu_graph_store.py b/tests/test_storage/test_kuzu_graph_store.py index b0b66abd..d0652e1b 100644 --- a/tests/test_storage/test_kuzu_graph_store.py +++ b/tests/test_storage/test_kuzu_graph_store.py @@ -1,9 +1,12 @@ -import pytest - from typing import Optional + import kuzu +import pytest + +from motleycrew.common import GraphStoreType from motleycrew.storage import MotleyGraphNode from motleycrew.storage import MotleyKuzuGraphStore +from tests.test_storage import GraphStoreFixtures class Entity(MotleyGraphNode): @@ -19,214 +22,15 @@ def database(tmpdir): return db -class TestMotleyKuzuGraphStore: +class TestMotleyKuzuGraphStore(GraphStoreFixtures): def test_set_get_node_id(self): entity = Entity(int_param=1) MotleyKuzuGraphStore._set_node_id(node=entity, node_id=2) assert entity.id == 2 - def test_insert_new_node(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity = Entity(int_param=1) - created_entity = graph_store.insert_node(entity) - assert created_entity.id is not None - assert entity.id is not None # mutated in place - - def test_insert_node_and_retrieve(self, database): - graph_store = MotleyKuzuGraphStore(database) - - entity = Entity( - int_param=1, optional_str_param="test", optional_list_str_param=["a", "b"] - ) - inserted_entity = graph_store.insert_node(entity) - assert inserted_entity.id is not None - - retrieved_entity = graph_store.get_node_by_class_and_id( - node_class=Entity, node_id=inserted_entity.id - ) - assert inserted_entity == retrieved_entity - - def test_insert_node_with_id_already_set(self, database): - graph_store = MotleyKuzuGraphStore(database) + @pytest.mark.parametrize("graph_store", [GraphStoreType.KUZU], indirect=True) + def test_insert_node_with_id_already_set(self, graph_store): entity = Entity(int_param=1) MotleyKuzuGraphStore._set_node_id(node=entity, node_id=2) with pytest.raises(AssertionError): graph_store.insert_node(entity) - - def test_check_node_exists_true(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity = Entity(int_param=1) - - graph_store.insert_node(entity) - assert graph_store.check_node_exists(entity) - - def test_check_node_exists_false(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity = Entity(int_param=1) - assert graph_store.check_node_exists(entity) is False - - MotleyKuzuGraphStore._set_node_id(node=entity, node_id=2) - assert graph_store.check_node_exists(entity) is False - - graph_store.ensure_node_table(Entity) - assert graph_store.check_node_exists(entity) is False - - def test_create_relation(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - graph_store.insert_node(entity1) - graph_store.insert_node(entity2) - - graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") - - assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - assert graph_store.check_relation_exists( - from_node=entity1, to_node=entity2, label="p" - ) - assert not graph_store.check_relation_exists( - from_node=entity1, to_node=entity2, label="q" - ) - assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) - - def test_upsert_triplet(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - graph_store.upsert_triplet(from_node=entity1, to_node=entity2, label="p") - - assert graph_store.check_node_exists(entity1) - assert graph_store.check_node_exists(entity2) - - assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - assert graph_store.check_relation_exists( - from_node=entity1, to_node=entity2, label="p" - ) - assert not graph_store.check_relation_exists( - from_node=entity1, to_node=entity2, label="q" - ) - assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) - - def test_nodes_do_not_exist(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - - assert not graph_store.check_node_exists(entity1) - assert not graph_store.check_node_exists(entity2) - - assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - assert not graph_store.check_relation_exists( - from_node=entity2, to_node=entity1, label="p" - ) - - def test_relation_does_not_exist(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - - assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) - - graph_store.insert_node(entity1) - graph_store.insert_node(entity2) - - assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1) - - def test_delete_node(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity = Entity(int_param=1) - graph_store.insert_node(entity) - assert graph_store.check_node_exists(entity) is True - - graph_store.delete_node(entity) - assert graph_store.check_node_exists(entity) is False - - entity.int_param = 2 # check that entity is not frozen - - def test_delete_entity_with_relations(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - - graph_store.insert_node(entity1) - graph_store.insert_node(entity2) - graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") - assert ( - graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - is True - ) - - graph_store.delete_node(entity1) - assert graph_store.check_node_exists(entity1) is False - assert graph_store.check_node_exists(entity2) is True - assert ( - graph_store.check_relation_exists(from_node=entity1, to_node=entity2) - is False - ) - - def test_set_property(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity = Entity(int_param=1) - graph_store.insert_node(entity) - assert entity.optional_str_param is None - assert ( - graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param - is None - ) - - entity.optional_str_param = "test" - assert ( - graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param - == "test" - ) - - entity.optional_list_str_param = ["a", "b"] - assert graph_store.get_node_by_class_and_id( - Entity, entity.id - ).optional_list_str_param == [ - "a", - "b", - ] - - def test_run_cypher_query(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2) - - graph_store.insert_node(entity1) - graph_store.insert_node(entity2) - graph_store.create_relation(from_node=entity1, to_node=entity2, label="p") - - query = """ - MATCH (a:Entity {int_param: 1})-[r]->(b:Entity {int_param: 2}) - RETURN a, r, b - """ - result = graph_store.run_cypher_query(query) - assert len(result) == 1 - assert len(result[0]) == 3 - - a, r, b = result[0] - assert a["int_param"] == 1 - assert b["int_param"] == 2 - assert r["_label"] == "p" - - def test_run_cypher_query_with_container(self, database): - graph_store = MotleyKuzuGraphStore(database) - entity1 = Entity(int_param=1) - entity2 = Entity(int_param=2, optional_list_str_param=["a", "b"]) - - graph_store.insert_node(entity1) - graph_store.insert_node(entity2) - - query = """ - MATCH (a:Entity) - WHERE a.int_param = 2 - RETURN a - """ - result = graph_store.run_cypher_query(query, container=Entity) - assert len(result) == 1 - assert isinstance(result[0], Entity) - assert result[0].int_param == 2 - assert result[0].optional_list_str_param == ["a", "b"]