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