from ims_mcp.clients.dataset import DatasetLookup


class _Dataset:
    def __init__(self, dataset_id, name):
        self.id = dataset_id
        self.name = name


class _Rag:
    def __init__(self):
        self.calls = 0

    def list_datasets(self, name=None, id=None):
        self.calls += 1
        if name == "known":
            return [_Dataset("d1", "known")]
        if id == "d1":
            return [_Dataset("d1", "known")]
        return []


def test_positive_and_negative_cache():
    rag = _Rag()
    lookup = DatasetLookup(ragflow=rag, ttl_seconds=300)

    assert lookup.get_id("known") == "d1"
    assert lookup.get_id("known") == "d1"
    assert lookup.get_id("missing") is None
    assert lookup.get_id("missing") is None

    assert rag.calls == 2


def test_invalidate_clears_cache():
    rag = _Rag()
    lookup = DatasetLookup(ragflow=rag, ttl_seconds=300)
    assert lookup.get_id("known") == "d1"
    lookup.invalidate()
    assert lookup.get_id("known") == "d1"
    assert rag.calls == 2
