Spaces:
Running
Running
| import os | |
| import unittest | |
| from unittest.mock import MagicMock, patch | |
| ENV_DEFAULTS = { | |
| "BRAINPAT_TOKEN": "test-token", | |
| "MODELS_MODE": "local", | |
| "EMBEDDINGS_LOCAL_MODEL": "local-model", | |
| "EMBEDDINGS_SMALL_MODEL": "small-model", | |
| "EMBEDDING_NODES_DIMENSION": "3", | |
| "EMBEDDING_TRIPLETS_DIMENSION": "3", | |
| "EMBEDDING_OBSERVATIONS_DIMENSION": "3", | |
| "EMBEDDING_DATA_DIMENSION": "3", | |
| "EMBEDDING_RELATIONSHIPS_DIMENSION": "3", | |
| "REDIS_HOST": "localhost", | |
| "REDIS_PORT": "6379", | |
| "NEO4J_HOST": "localhost", | |
| "NEO4J_PORT": "7687", | |
| "NEO4J_USERNAME": "neo4j", | |
| "NEO4J_PASSWORD": "password", | |
| "MILVUS_HOST": "localhost", | |
| "MILVUS_PORT": "19530", | |
| "MONGO_CONNECTION_STRING": "mongodb://localhost:27017", | |
| "CELERY_WORKER_CONCURRENCY": "1", | |
| "OLLAMA_HOST": "localhost", | |
| "OLLAMA_PORT": "11434", | |
| "OLLAMA_LLM_SMALL_MODEL": "small", | |
| "OLLAMA_LLM_LARGE_MODEL": "large", | |
| "PIPELINE_MODE": "accurate", | |
| } | |
| for key, value in ENV_DEFAULTS.items(): | |
| os.environ.setdefault(key, value) | |
| class NetworkXUpsertContractTests(unittest.TestCase): | |
| def test_same_uuid_upsert_keeps_one_node_and_preserves_uuid(self): | |
| from src.lib.postgresql.graph_store import _BrainGraph | |
| brain = _BrainGraph("brain-a") | |
| first = brain.upsert_node( | |
| ["PERSON"], | |
| {"uuid": "alice-1", "name": "Alice"}, | |
| {"uuid": "alice-1", "name": "Alice", "description": "v1"}, | |
| ) | |
| second = brain.upsert_node( | |
| ["PERSON"], | |
| {"uuid": "alice-1", "name": "Alice"}, | |
| {"uuid": "alice-1", "name": "Alice", "description": "v2"}, | |
| ) | |
| self.assertEqual(first, "alice-1") | |
| self.assertEqual(second, "alice-1") | |
| self.assertEqual(len(list(brain.graph.nodes)), 1) | |
| self.assertEqual(brain.node_data("alice-1")["uuid"], "alice-1") | |
| self.assertEqual(brain.node_data("alice-1")["description"], "v1 | v2") | |
| def test_same_name_different_uuid_does_not_overwrite_existing_uuid(self): | |
| from src.lib.postgresql.graph_store import _BrainGraph | |
| brain = _BrainGraph("brain-a") | |
| brain.upsert_node( | |
| ["PERSON"], | |
| {"uuid": "alice-1", "name": "Alice"}, | |
| {"uuid": "alice-1", "name": "Alice"}, | |
| ) | |
| brain.upsert_node( | |
| ["PERSON"], | |
| {"uuid": "alice-2", "name": "Alice"}, | |
| {"uuid": "alice-2", "name": "Alice"}, | |
| ) | |
| self.assertEqual(len(list(brain.graph.nodes)), 2) | |
| self.assertEqual(brain.node_data("alice-1")["uuid"], "alice-1") | |
| self.assertEqual(brain.node_data("alice-2")["uuid"], "alice-2") | |
| def test_relationship_upsert_by_uuid_is_idempotent(self): | |
| from src.lib.postgresql.graph_store import PostgreSQLGraphStore | |
| store = PostgreSQLGraphStore.__new__(PostgreSQLGraphStore) | |
| from src.lib.postgresql.graph_store import _BrainGraph | |
| brain = _BrainGraph("brain-a") | |
| brain.upsert_node(["PERSON"], {"uuid": "a"}, {"uuid": "a", "name": "Alice"}) | |
| brain.upsert_node(["PRODUCT"], {"uuid": "b"}, {"uuid": "b", "name": "Widget"}) | |
| store._brains = {"brain-a": brain} | |
| store._brains_lock = __import__("threading").RLock() | |
| store._persist_relationship = MagicMock() | |
| store._persist_node = MagicMock() | |
| first = store.merge_relationship( | |
| "brain-a", | |
| ["PERSON"], | |
| "Alice", | |
| ["PRODUCT"], | |
| "Widget", | |
| "OWNS", | |
| {"uuid": "rel-1", "description": "v1"}, | |
| subject_uuid="a", | |
| object_uuid="b", | |
| ) | |
| second = store.merge_relationship( | |
| "brain-a", | |
| ["PERSON"], | |
| "Alice", | |
| ["PRODUCT"], | |
| "Widget", | |
| "OWNS", | |
| {"uuid": "rel-1", "description": "v2"}, | |
| subject_uuid="a", | |
| object_uuid="b", | |
| ) | |
| self.assertIsNotNone(first) | |
| self.assertIsNotNone(second) | |
| edges = list(brain.graph.edges(keys=True, data=True)) | |
| self.assertEqual(len(edges), 1) | |
| self.assertEqual(edges[0][2], "rel-1") | |
| self.assertEqual(edges[0][3]["description"], "v2") | |
| def test_relationship_upsert_creates_missing_endpoint_nodes(self): | |
| from src.lib.postgresql.graph_store import PostgreSQLGraphStore | |
| from src.lib.postgresql.graph_store import _BrainGraph | |
| store = PostgreSQLGraphStore.__new__(PostgreSQLGraphStore) | |
| brain = _BrainGraph("brain-a") | |
| store._brains = {"brain-a": brain} | |
| store._brains_lock = __import__("threading").RLock() | |
| store._persist_node = MagicMock() | |
| store._persist_relationship = MagicMock() | |
| result = store.merge_relationship( | |
| "brain-a", | |
| ["PERSON"], | |
| "Alice", | |
| ["PRODUCT"], | |
| "Widget", | |
| "OWNS", | |
| {"uuid": "rel-1", "description": "created"}, | |
| subject_uuid="a", | |
| object_uuid="b", | |
| ) | |
| self.assertIsNotNone(result) | |
| self.assertEqual(set(brain.graph.nodes), {"a", "b"}) | |
| self.assertEqual(len(list(brain.graph.edges(keys=True))), 1) | |
| store._persist_node.assert_called() | |
| store._persist_relationship.assert_called_once() | |
| def test_normalize_read_sql_rewrites_jsonb_ilike(self): | |
| from src.lib.postgresql.graph_store import normalize_read_sql | |
| sql = ( | |
| "SELECT uuid, data FROM kg_nodes WHERE data->>'name' ILIKE '%Reminding%' " | |
| "OR data->'description' ILIKE '%remind%' LIMIT 50" | |
| ) | |
| fixed = normalize_read_sql(sql) | |
| self.assertIn("data->>'description' ILIKE", fixed) | |
| self.assertNotIn("data->'description' ILIKE", fixed) | |
| def test_stale_cache_reloads_when_db_counts_diverge(self): | |
| from src.lib.postgresql.graph_store import PostgreSQLGraphStore | |
| from src.lib.postgresql.graph_store import _BrainGraph | |
| store = PostgreSQLGraphStore.__new__(PostgreSQLGraphStore) | |
| stale = _BrainGraph("brain-a") | |
| store._brains = {"brain-a": stale} | |
| store._brains_lock = __import__("threading").RLock() | |
| store._db_graph_counts = MagicMock(return_value=(2, 1)) | |
| fresh = _BrainGraph("brain-a") | |
| fresh.graph.add_node("a", name="Alice", labels=["PERSON"]) | |
| fresh.graph.add_node("b", name="Bob", labels=["PERSON"]) | |
| fresh.graph.add_edge("a", "b", key="r1", rel_type="KNOWS") | |
| store._reload_brain_locked = MagicMock(return_value=fresh) | |
| with patch("src.lib.postgresql.graph_store.require_registered_brain"): | |
| loaded = store._load_brain("brain-a") | |
| self.assertIs(loaded, fresh) | |
| store._reload_brain_locked.assert_called_once_with("brain-a") | |
| class Neo4jUpsertContractTests(unittest.TestCase): | |
| def test_add_nodes_merges_by_uuid_not_name(self): | |
| from src.constants.kg import Node | |
| from src.lib.neo4j.client import Neo4jClient | |
| client = Neo4jClient.__new__(Neo4jClient) | |
| client.ensure_database = MagicMock() | |
| queries = [] | |
| def _capture(query, brain_id): | |
| queries.append(query) | |
| return "ok" | |
| client._execute_query_with_retry = _capture | |
| client._clean_labels = lambda labels: [str(l).upper() for l in labels] | |
| client._format_property_key = lambda key: key | |
| client._format_value = lambda value: ( | |
| "null" | |
| if value is None | |
| else f"'{value}'" | |
| if isinstance(value, str) | |
| else str(value) | |
| ) | |
| client._clean_property_key = lambda key: key | |
| node = Node( | |
| uuid="alice-1", | |
| name="Alice", | |
| labels=["PERSON"], | |
| description="v1", | |
| properties={}, | |
| ) | |
| client.add_nodes([node], "brain-a") | |
| self.assertEqual(len(queries), 1) | |
| self.assertIn("MERGE (n:PERSON {uuid: 'alice-1'})", queries[0]) | |
| self.assertNotIn("{name:", queries[0]) | |
| def test_add_relationship_merges_by_endpoint_and_relationship_uuid(self): | |
| from src.constants.kg import Node, Predicate | |
| from src.lib.neo4j.client import Neo4jClient | |
| client = Neo4jClient.__new__(Neo4jClient) | |
| client.ensure_database = MagicMock() | |
| client.driver = MagicMock() | |
| client._clean_labels = lambda labels: [str(l).upper() for l in labels] | |
| client._format_value = lambda value: ( | |
| "null" | |
| if value is None | |
| else f"'{value}'" | |
| if isinstance(value, str) | |
| else str(value) | |
| ) | |
| subject = Node(uuid="a", name="Alice", labels=["PERSON"], properties={}) | |
| obj = Node(uuid="b", name="Widget", labels=["PRODUCT"], properties={}) | |
| predicate = Predicate( | |
| uuid="rel-1", | |
| name="OWNS", | |
| description="owns", | |
| properties={"v_id": "v1"}, | |
| flow_key="flow-1", | |
| ) | |
| client.add_relationship(subject, predicate, obj, "brain-a") | |
| query = client.driver.execute_query.call_args.args[0] | |
| self.assertIn("WHERE a['uuid'] = 'a'", query) | |
| self.assertIn("WHERE b['uuid'] = 'b'", query) | |
| self.assertIn("MERGE (a)-[r:OWNS {uuid: 'rel-1'}]->(b)", query) | |
| if __name__ == "__main__": | |
| unittest.main() | |