langchain-mongodb: add unit tests for MongoDBChatMessageHistory

libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py

@@ -31,7 +31,7 @@ def test_memory_with_message_store() -> None:
     assert "This is me, the AI" in messages_json
     assert "This is me, the human" in messages_json
 
-    # remove the record from Azure Cosmos DB, so the next test run won't pick it up
+    # remove the record from MongoDB, so the next test run won't pick it up
     memory.chat_memory.clear()
 
     assert memory.chat_memory.messages == []

libs/partners/mongodb/tests/unit_tests/test_chat_message_histories.py

@@ -0,0 +1,40 @@
+import json
+
+from langchain.memory import ConversationBufferMemory
+from langchain_core.messages import message_to_dict
+
+from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
+from tests.utils import MockCollection
+
+
+class PatchedMongoDBChatMessageHistory(MongoDBChatMessageHistory):
+    def __init__(self) -> None:
+        self.session_id = "test-session"
+        self.database_name = "test-database"
+        self.collection_name = "test-collection"
+        self.collection = MockCollection()
+
+
+def test_memory_with_message_store() -> None:
+    """Test the memory with a message store."""
+    # setup MongoDB as a message store
+    message_history = PatchedMongoDBChatMessageHistory()
+    memory = ConversationBufferMemory(
+        memory_key="baz", chat_memory=message_history, return_messages=True
+    )
+
+    # add some messages
+    memory.chat_memory.add_ai_message("This is me, the AI")
+    memory.chat_memory.add_user_message("This is me, the human")
+
+    # get the message history from the memory store and turn it into a json
+    messages = memory.chat_memory.messages
+    messages_json = json.dumps([message_to_dict(msg) for msg in messages])
+
+    assert "This is me, the AI" in messages_json
+    assert "This is me, the human" in messages_json
+
+    # remove the record from MongoDB, so the next test run won't pick it up
+    memory.chat_memory.clear()
+
+    assert memory.chat_memory.messages == []

libs/partners/mongodb/tests/utils.py

@@ -169,18 +169,21 @@ def insert_many(self, to_insert: List[Any], *args, **kwargs) -> InsertManyResult
             [k["_id"] for k in mongodb_inserts], acknowledged=True
         )
 
+    def insert_one(self, to_insert: Any, *args, **kwargs) -> Any:  # type: ignore
+        return self.insert_many([to_insert])
+
     def find_one(self, find_query: Dict[str, Any]) -> Optional[Dict[str, Any]]:  # type: ignore
+        find = self.find(find_query) or [None]  # type: ignore
+        return find[0]
+
+    def find(self, find_query: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:  # type: ignore
         def _is_match(item: Dict[str, Any]) -> bool:
             for key, match_val in find_query.items():
                 if item.get(key) != match_val:
                     return False
             return True
 
-        # Return the first element to match
-        for document in self._data:
-            if _is_match(document):
-                return document
-        return None
+        return [document for document in self._data if _is_match(document)]
 
     def update_one(  # type: ignore
         self,