def mongo_chat_store(
    mongo_container: Dict[str, Union[str, Container]],
) -> Generator[MongoChatStore, None, None]:
    """Fixture to create a MongoChatStore instance connected to the test container."""
    chat_store = None
    try:
        chat_store = MongoChatStore(
            mongo_uri=mongo_container["mongodb_uri"],
            db_name="test_db",
            collection_name="test_chats",
        )
        yield chat_store
    finally:
        if chat_store and hasattr(chat_store, "_collection"):
            # Clean up by dropping the collection
            chat_store._collection.drop()
Fixture to create a MongoChatStore instance connected to the test container.
llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-mongo/tests/test_chat_store_mongo_chat_store.py