import pytest

from algebraic_auction_studio_for_robotic_fle.dsl import Agent, Task
from algebraic_auction_studio_for_robotic_fle.engine_admm import admm_solve


def test_admm_solve_basic_allocation():
    # Two agents with differing costs
    a1 = Agent(id="robot_A", capabilities=["lift"], cost_model={"weight": 1.0})
    a2 = Agent(id="robot_B", capabilities=["navigate"], cost_model={"weight": 0.8})

    # Two simple tasks with values and requirements
    t1 = Task(id="task_1", requirements={"duration": 3, "energy": 2}, deadline=100.0, value=10.0)
    t2 = Task(id="task_2", requirements={"duration": 4, "energy": 3}, deadline=120.0, value=12.0)

    alloc = admm_solve([a1, a2], [t1, t2])

    # Expect an allocation for both tasks to some agent
    assert set(alloc.keys()) == {t1.id, t2.id}
    assert alloc[t1.id] in {a1.id, a2.id}
    assert alloc[t2.id] in {a1.id, a2.id}