unit tests

tests/utils/maintenance/test_temporal_operations.py

@@ -33,9 +33,9 @@ def create_test_data():
     now = datetime.now()
 
     # Create nodes
-    node1 = EntityNode(uuid='1', name='Node1', labels=['Person'], created_at=now)
-    node2 = EntityNode(uuid='2', name='Node2', labels=['Person'], created_at=now)
-    node3 = EntityNode(uuid='3', name='Node3', labels=['Person'], created_at=now)
+    node1 = EntityNode(uuid='1', name='Node1', labels=['Person'], created_at=now, group_id='1')
+    node2 = EntityNode(uuid='2', name='Node2', labels=['Person'], created_at=now, group_id='1')
+    node3 = EntityNode(uuid='3', name='Node3', labels=['Person'], created_at=now, group_id='1')
 
     # Create edges
     existing_edge1 = EntityEdge(
@@ -45,6 +45,7 @@ def create_test_data():
         name='KNOWS',
         fact='Node1 knows Node2',
         created_at=now,
+        group_id='1',
     )
     existing_edge2 = EntityEdge(
         uuid='e2',
@@ -53,6 +54,7 @@ def create_test_data():
         name='LIKES',
         fact='Node2 likes Node3',
         created_at=now,
+        group_id='1',
     )
     new_edge1 = EntityEdge(
         uuid='e3',
@@ -61,6 +63,7 @@ def create_test_data():
         name='WORKS_WITH',
         fact='Node1 works with Node3',
         created_at=now,
+        group_id='1',
     )
     new_edge2 = EntityEdge(
         uuid='e4',
@@ -69,6 +72,7 @@ def create_test_data():
         name='DISLIKES',
         fact='Node1 dislikes Node2',
         created_at=now,
+        group_id='1',
     )
 
     return {
@@ -258,8 +262,8 @@ def test_prepare_invalidation_context_sorting():
     now = datetime.now()
 
     # Create nodes
-    node1 = EntityNode(uuid='1', name='Node1', labels=['Person'], created_at=now)
-    node2 = EntityNode(uuid='2', name='Node2', labels=['Person'], created_at=now)
+    node1 = EntityNode(uuid='1', name='Node1', labels=['Person'], created_at=now, group_id='1')
+    node2 = EntityNode(uuid='2', name='Node2', labels=['Person'], created_at=now, group_id='1')
 
     # Create edges with different timestamps
     edge1 = EntityEdge(

tests/utils/search/search_utils_test.py

@@ -19,12 +19,12 @@ async def test_hybrid_node_search_deduplication():
     ) as mock_similarity_search:
         # Set up mock return values
         mock_fulltext_search.side_effect = [
-            [EntityNode(uuid='1', name='Alice', labels=['Entity'])],
-            [EntityNode(uuid='2', name='Bob', labels=['Entity'])],
+            [EntityNode(uuid='1', name='Alice', labels=['Entity'], group_id='1')],
+            [EntityNode(uuid='2', name='Bob', labels=['Entity'], group_id='1')],
         ]
         mock_similarity_search.side_effect = [
-            [EntityNode(uuid='1', name='Alice', labels=['Entity'])],
-            [EntityNode(uuid='3', name='Charlie', labels=['Entity'])],
+            [EntityNode(uuid='1', name='Alice', labels=['Entity'], group_id='1')],
+            [EntityNode(uuid='3', name='Charlie', labels=['Entity'], group_id='1')],
         ]
 
         # Call the function with test data
@@ -120,8 +120,8 @@ async def test_hybrid_node_search_with_limit():
         assert mock_fulltext_search.call_count == 1
         assert mock_similarity_search.call_count == 1
         # Verify that the limit was passed to the search functions
-        mock_fulltext_search.assert_called_with('Test', mock_driver, 2)
-        mock_similarity_search.assert_called_with([0.1, 0.2, 0.3], mock_driver, 2)
+        mock_fulltext_search.assert_called_with('Test', mock_driver, ['1'], 2)
+        mock_similarity_search.assert_called_with([0.1, 0.2, 0.3], mock_driver, ['1'], 2)
 
 
 @pytest.mark.asyncio
@@ -155,5 +155,5 @@ async def test_hybrid_node_search_with_limit_and_duplicates():
         assert set(node.name for node in results) == {'Alice', 'Bob', 'Charlie'}
         assert mock_fulltext_search.call_count == 1
         assert mock_similarity_search.call_count == 1
-        mock_fulltext_search.assert_called_with('Test', mock_driver, 4)
-        mock_similarity_search.assert_called_with([0.1, 0.2, 0.3], mock_driver, 4)
+        mock_fulltext_search.assert_called_with('Test', mock_driver, ['1'], 4)
+        mock_similarity_search.assert_called_with([0.1, 0.2, 0.3], mock_driver, ['1'], 4)