37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
from typing import Any, Dict
|
|
from unittest.mock import patch
|
|
|
|
from great_ai import Trace
|
|
from great_ai.persistence.mongodb_driver import MongoDbDriver
|
|
from pymongo import MongoClient
|
|
from pymongo.synchronous.collection import Collection
|
|
|
|
|
|
def test_mongo_update_uses_valid_replacement() -> None:
|
|
trace = Trace(
|
|
trace_id="id",
|
|
created="2026-01-01",
|
|
original_execution_time_ms=0,
|
|
logged_values={},
|
|
models=[],
|
|
tags=[],
|
|
feedback="correct",
|
|
)
|
|
driver = MongoDbDriver.__new__(MongoDbDriver)
|
|
driver.mongo_connection_string = "mongodb://127.0.0.1:1"
|
|
driver.mongo_database = "test"
|
|
client: MongoClient[Dict[str, Any]] = MongoClient(
|
|
driver.mongo_connection_string, connect=False
|
|
)
|
|
with (
|
|
patch("great_ai.persistence.mongodb_driver.MongoClient") as factory,
|
|
patch.object(
|
|
Collection, "_update_retryable", return_value={"n": 1, "nModified": 1}
|
|
) as update,
|
|
):
|
|
factory.__getitem__.return_value.return_value = client
|
|
driver.update("id", trace)
|
|
assert update.call_args.args[0] == {"_id": "id"}
|
|
document = update.call_args.args[1]
|
|
assert document["_id"] == "id"
|
|
assert Trace.model_validate(document) == trace
|