You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
langchain/libs/community/tests/integration_tests/indexes/test_document_manager.py

269 lines
9.0 KiB
Python

from datetime import datetime
from unittest.mock import patch
import pytest
import pytest_asyncio
from langchain_community.indexes._document_manager import MongoDocumentManager
@pytest.fixture
@pytest.mark.requires("pymongo")
def manager() -> MongoDocumentManager:
"""Initialize the test MongoDB and yield the DocumentManager instance."""
document_manager = MongoDocumentManager(
namespace="kittens",
mongodb_url="mongodb://langchain:langchain@localhost:6022/",
db_name="test_db",
collection_name="test_collection",
)
return document_manager
@pytest_asyncio.fixture
@pytest.mark.requires("motor")
async def amanager() -> MongoDocumentManager:
"""Initialize the test MongoDB and yield the DocumentManager instance."""
document_manager = MongoDocumentManager(
namespace="kittens",
mongodb_url="mongodb://langchain:langchain@localhost:6022/",
db_name="test_db",
collection_name="test_collection",
)
return document_manager
@pytest.mark.requires("pymongo")
def test_update(manager: MongoDocumentManager) -> None:
"""Test updating records in the MongoDB."""
read_keys = manager.list_keys()
updated_keys = ["update_key1", "update_key2", "update_key3"]
manager.update(updated_keys)
all_keys = manager.list_keys()
assert sorted(all_keys) == sorted(read_keys + updated_keys)
@pytest.mark.asyncio
@pytest.mark.requires("motor")
async def test_aupdate(amanager: MongoDocumentManager) -> None:
"""Test updating records in the MongoDB."""
read_keys = await amanager.alist_keys()
aupdated_keys = ["aupdate_key1", "aupdate_key2", "aupdate_key3"]
await amanager.aupdate(aupdated_keys)
all_keys = await amanager.alist_keys()
assert sorted(all_keys) == sorted(read_keys + aupdated_keys)
@pytest.mark.requires("pymongo")
def test_update_timestamp(manager: MongoDocumentManager) -> None:
"""Test updating records with timestamps in MongoDB."""
with patch.object(
manager, "get_time", return_value=datetime(2024, 2, 23).timestamp()
):
manager.update(["key1"])
records = list(
manager.sync_collection.find({"namespace": manager.namespace, "key": "key1"})
)
assert [
{
"key": record["key"],
"namespace": record["namespace"],
"updated_at": record["updated_at"],
"group_id": record.get("group_id"),
}
for record in records
] == [
{
"group_id": None,
"key": "key1",
"namespace": "kittens",
"updated_at": datetime(2024, 2, 23).timestamp(),
}
]
@pytest.mark.requires("motor")
async def test_aupdate_timestamp(amanager: MongoDocumentManager) -> None:
"""Test asynchronously updating records with timestamps in MongoDB."""
with patch.object(
amanager, "aget_time", return_value=datetime(2024, 2, 23).timestamp()
):
await amanager.aupdate(["key1"])
records = [
doc
async for doc in amanager.async_collection.find(
{"namespace": amanager.namespace, "key": "key1"}
)
]
assert [
{
"key": record["key"],
"namespace": record["namespace"],
"updated_at": record["updated_at"],
"group_id": record.get("group_id"),
}
for record in records
] == [
{
"group_id": None,
"key": "key1",
"namespace": "kittens",
"updated_at": datetime(2024, 2, 23).timestamp(),
}
]
@pytest.mark.requires("pymongo")
def test_exists(manager: MongoDocumentManager) -> None:
"""Test checking if keys exist in MongoDB."""
keys = ["key1", "key2", "key3"]
manager.update(keys)
exists = manager.exists(keys)
assert len(exists) == len(keys)
assert all(exists)
exists = manager.exists(["key1", "key4"])
assert len(exists) == 2
assert exists == [True, False]
@pytest.mark.requires("motor")
async def test_aexists(amanager: MongoDocumentManager) -> None:
"""Test asynchronously checking if keys exist in MongoDB."""
keys = ["key1", "key2", "key3"]
await amanager.aupdate(keys)
exists = await amanager.aexists(keys)
assert len(exists) == len(keys)
assert all(exists)
exists = await amanager.aexists(["key1", "key4"])
assert len(exists) == 2
assert exists == [True, False]
@pytest.mark.requires("pymongo")
def test_list_keys(manager: MongoDocumentManager) -> None:
"""Test listing keys in MongoDB."""
manager.delete_keys(manager.list_keys())
with patch.object(
manager, "get_time", return_value=datetime(2021, 1, 1).timestamp()
):
manager.update(["key1"])
with patch.object(
manager, "get_time", return_value=datetime(2022, 1, 1).timestamp()
):
manager.update(["key2"])
with patch.object(
manager, "get_time", return_value=datetime(2023, 1, 1).timestamp()
):
manager.update(["key3"])
with patch.object(
manager, "get_time", return_value=datetime(2024, 1, 1).timestamp()
):
manager.update(["key4"], group_ids=["group1"])
assert sorted(manager.list_keys()) == sorted(["key1", "key2", "key3", "key4"])
assert sorted(manager.list_keys(after=datetime(2022, 2, 1).timestamp())) == sorted(
["key3", "key4"]
)
assert sorted(manager.list_keys(group_ids=["group1", "group2"])) == sorted(["key4"])
@pytest.mark.requires("motor")
async def test_alist_keys(amanager: MongoDocumentManager) -> None:
"""Test asynchronously listing keys in MongoDB."""
await amanager.adelete_keys(await amanager.alist_keys())
with patch.object(
amanager, "aget_time", return_value=datetime(2021, 1, 1).timestamp()
):
await amanager.aupdate(["key1"])
with patch.object(
amanager, "aget_time", return_value=datetime(2022, 1, 1).timestamp()
):
await amanager.aupdate(["key2"])
with patch.object(
amanager, "aget_time", return_value=datetime(2023, 1, 1).timestamp()
):
await amanager.aupdate(["key3"])
with patch.object(
amanager, "aget_time", return_value=datetime(2024, 1, 1).timestamp()
):
await amanager.aupdate(["key4"], group_ids=["group1"])
assert sorted(await amanager.alist_keys()) == sorted(
["key1", "key2", "key3", "key4"]
)
assert sorted(
await amanager.alist_keys(after=datetime(2022, 2, 1).timestamp())
) == sorted(["key3", "key4"])
assert sorted(await amanager.alist_keys(group_ids=["group1", "group2"])) == sorted(
["key4"]
)
@pytest.mark.requires("pymongo")
def test_namespace_is_used(manager: MongoDocumentManager) -> None:
"""Verify that namespace is taken into account for all operations in MongoDB."""
manager.delete_keys(manager.list_keys())
manager.update(["key1", "key2"], group_ids=["group1", "group2"])
manager.sync_collection.insert_many(
[
{"key": "key1", "namespace": "puppies", "group_id": None},
{"key": "key3", "namespace": "puppies", "group_id": None},
]
)
assert sorted(manager.list_keys()) == sorted(["key1", "key2"])
manager.delete_keys(["key1"])
assert sorted(manager.list_keys()) == sorted(["key2"])
manager.update(["key3"], group_ids=["group3"])
assert (
manager.sync_collection.find_one({"key": "key3", "namespace": "kittens"})[
"group_id"
]
== "group3"
)
@pytest.mark.requires("motor")
async def test_anamespace_is_used(amanager: MongoDocumentManager) -> None:
"""
Verify that namespace is taken into account for all operations
in MongoDB asynchronously.
"""
await amanager.adelete_keys(await amanager.alist_keys())
await amanager.aupdate(["key1", "key2"], group_ids=["group1", "group2"])
await amanager.async_collection.insert_many(
[
{"key": "key1", "namespace": "puppies", "group_id": None},
{"key": "key3", "namespace": "puppies", "group_id": None},
]
)
assert sorted(await amanager.alist_keys()) == sorted(["key1", "key2"])
await amanager.adelete_keys(["key1"])
assert sorted(await amanager.alist_keys()) == sorted(["key2"])
await amanager.aupdate(["key3"], group_ids=["group3"])
assert (
await amanager.async_collection.find_one(
{"key": "key3", "namespace": "kittens"}
)
)["group_id"] == "group3"
@pytest.mark.requires("pymongo")
def test_delete_keys(manager: MongoDocumentManager) -> None:
"""Test deleting keys from MongoDB."""
manager.update(["key1", "key2", "key3"])
manager.delete_keys(["key1", "key2"])
remaining_keys = manager.list_keys()
assert sorted(remaining_keys) == sorted(["key3"])
@pytest.mark.requires("motor")
async def test_adelete_keys(amanager: MongoDocumentManager) -> None:
"""Test asynchronously deleting keys from MongoDB."""
await amanager.aupdate(["key1", "key2", "key3"])
await amanager.adelete_keys(["key1", "key2"])
remaining_keys = await amanager.alist_keys()
assert sorted(remaining_keys) == sorted(["key3"])