from catopt_query.adapters.postgres_adapter import map_postgres_to_canonical
from catopt_query.adapters.mongo_adapter import map_mongo_to_canonical
from catopt_query.protocol import LocalProblem
from catopt_query.core import CanonicalPlan


def test_postgres_adapter_maps_to_canonical():
    lp = LocalProblem(
        shard_id="pg-1",
        projected_attributes=["a"],
        predicates=["a > 1"],
        costs={"cpu": 1.0},
    )
    can = map_postgres_to_canonical(lp)
    assert isinstance(can, CanonicalPlan)
    assert can.total_cost == 1.0


def test_mongo_adapter_maps_to_canonical():
    lp = LocalProblem(
        shard_id="mongo-1",
        projected_attributes=["b"],
        predicates=["b != NULL"],
        costs={"io": 0.2},
    )
    can = map_mongo_to_canonical(lp)
    assert isinstance(can, CanonicalPlan)
    assert can.total_cost == 0.2