tests: add tests for tree memory add

Author: CaralHsi

src/memos/memories/textual/tree_text_memory/organize/manager.py

@@ -212,7 +212,9 @@ def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> s
         merged_background = f"{original_meta.background}\n⟵MERGED⟶\n{source_meta.background}"
         merged_embedding = self.embedder.embed([merged_text])[0]
 
-        merged_confidence = float((original_meta.confidence + source_meta.confidence) / 2)
+        original_conf = original_meta.confidence or 0.0
+        source_conf = source_meta.confidence or 0.0
+        merged_confidence = float((original_conf + source_conf) / 2)
         merged_usage = list(set((original_meta.usage or []) + (source_meta.usage or [])))
 
         # Create new merged node

tests/memories/textual/test_tree.py

@@ -138,3 +138,27 @@ def test_drop_creates_backup_and_cleans(mock_tree_text_memory):
     mock_tree_text_memory.dump.assert_called_once()
     mock_tree_text_memory._cleanup_old_backups.assert_called_once()
     mock_tree_text_memory.graph_store.drop_database.assert_called_once()
+
+
+def test_add_returns_ids(mock_tree_text_memory):
+    # Mock the memory_manager.add to return specific IDs
+    dummy_ids = ["id1", "id2"]
+    mock_tree_text_memory.memory_manager.add = MagicMock(return_value=dummy_ids)
+
+    mock_items = [
+        TextualMemoryItem(
+            id="id1",
+            memory="Memory 1",
+            metadata=TreeNodeTextualMemoryMetadata(updated_at=None),
+        ),
+        TextualMemoryItem(
+            id="id2",
+            memory="Memory 2",
+            metadata=TreeNodeTextualMemoryMetadata(updated_at=None),
+        ),
+    ]
+
+    result = mock_tree_text_memory.add(mock_items)
+
+    assert result == dummy_ids
+    mock_tree_text_memory.memory_manager.add.assert_called_once_with(mock_items)

tests/memories/textual/test_tree_manager.py

@@ -145,3 +145,14 @@ def test_ensure_structure_path_reuses_existing(memory_manager, mock_graph_store)
     meta = TreeNodeTextualMemoryMetadata(key="hobby")
     node_id = memory_manager._ensure_structure_path("UserMemory", meta)
     assert node_id == "existing_node_id"
+
+
+def test_add_returns_written_node_ids(memory_manager):
+    memory = TextualMemoryItem(
+        memory="test memory",
+        metadata=TreeNodeTextualMemoryMetadata(embedding=[0.1] * 5, memory_type="UserMemory"),
+    )
+    ids = memory_manager.add([memory])
+    assert isinstance(ids, list)
+    assert all(isinstance(i, str) for i in ids)
+    assert len(ids) > 0