tests/test_constraints.py

"""Tests for the constraint engine."""

import pytest
from qsvaps.models import (
    Plan,
    PlanAction,
    PlanConstraint,
    ResourceConstraint,
    ConstraintType,
)
from qsvaps.constraint_engine import ConstraintEngine


def _simple_plan():
    """Two actions with a dependency."""
    return Plan(
        name="Simple",
        actions=[
            PlanAction(name="a", can_fail=True),
            PlanAction(name="b", can_fail=True),
        ],
        dependencies=[("a", "b")],
    )


def _resource_plan():
    """Two actions sharing a rate-limited resource."""
    return Plan(
        name="Resource",
        actions=[
            PlanAction(name="x", resources=["api"]),
            PlanAction(name="y", resources=["api"]),
        ],
        resource_constraints=[
            ResourceConstraint("api", max_concurrent=1),
        ],
    )


def _must_succeed_plan():
    """One action that must succeed."""
    return Plan(
        name="Must Succeed",
        actions=[
            PlanAction(name="critical", can_fail=False),
        ],
    )


class TestVariableAssignment:
    def test_success_variables(self):
        engine = ConstraintEngine(_simple_plan())
        assert "s_a" in engine.var_mapping.variables
        assert "s_b" in engine.var_mapping.variables
        assert engine.var_mapping.num_variables == 2

    def test_parallel_variables(self):
        engine = ConstraintEngine(_resource_plan())
        # s_x, s_y, p_x_y
        assert engine.var_mapping.num_variables == 3
        assert "p_x_y" in engine.var_mapping.variables


class TestDependencyConstraints:
    def test_generates_constraint(self):
        engine = ConstraintEngine(_simple_plan())
        constraints = engine.extract_constraints()
        dep_constraints = [
            c for c in constraints
            if c.constraint_type == ConstraintType.DEPENDENCY
        ]
        assert len(dep_constraints) == 1
        assert "requires" in dep_constraints[0].description.lower() or \
               "succeed" in dep_constraints[0].description.lower()

    def test_dependency_logic(self):
        """b succeeding without a should violate the dependency."""
        engine = ConstraintEngine(_simple_plan())
        constraints = engine.extract_constraints()

        # State: a=0, b=1 → b succeeds but a didn't → violation
        # x0=s_a=0, x1=s_b=1 → state = 0b10 = 2
        is_viol, violated = engine.evaluate_state(2, constraints)
        assert is_viol

        # State: a=1, b=1 → both succeed → no violation
        is_viol, violated = engine.evaluate_state(3, constraints)
        assert not is_viol


class TestResourceConstraints:
    def test_generates_constraint(self):
        engine = ConstraintEngine(_resource_plan())
        constraints = engine.extract_constraints()
        res_constraints = [
            c for c in constraints
            if c.constraint_type == ConstraintType.RESOURCE
        ]
        assert len(res_constraints) == 1

    def test_parallel_conflict(self):
        """Both actions succeed + run in parallel → violation."""
        engine = ConstraintEngine(_resource_plan())
        constraints = engine.extract_constraints()

        # Variables: s_x=0, s_y=1, p_x_y=2
        # State: s_x=1, s_y=1, p_x_y=1 → 0b111 = 7
        is_viol, _ = engine.evaluate_state(7, constraints)
        assert is_viol

        # State: s_x=1, s_y=1, p_x_y=0 → not parallel → OK
        is_viol, _ = engine.evaluate_state(3, constraints)
        assert not is_viol


class TestCompletionConstraints:
    def test_must_succeed(self):
        engine = ConstraintEngine(_must_succeed_plan())
        constraints = engine.extract_constraints()

        comp = [
            c for c in constraints
            if c.constraint_type == ConstraintType.COMPLETION
        ]
        assert len(comp) == 1

        # State 0: critical fails → violation
        is_viol, _ = engine.evaluate_state(0, constraints)
        assert is_viol

        # State 1: critical succeeds → OK
        is_viol, _ = engine.evaluate_state(1, constraints)
        assert not is_viol


class TestFindAllViolations:
    def test_simple_plan_violations(self):
        engine = ConstraintEngine(_simple_plan())
        constraints = engine.extract_constraints()
        violations = engine.find_all_violations(constraints)

        # 2 qubits → 4 states. The only violation is state 2 (b=1, a=0)
        assert 2 in violations
        assert 0 not in violations  # Both fail — no dependency violated
        assert 1 not in violations  # a=1, b=0 — OK
        assert 3 not in violations  # Both succeed — OK


class TestBooleanEvaluation:
    def test_simple_expressions(self):
        assert ConstraintEngine._eval_bool_expr(
            "x0 or x1", {"x0": True, "x1": False}
        ) is True
        assert ConstraintEngine._eval_bool_expr(
            "x0 and x1", {"x0": True, "x1": False}
        ) is False
        assert ConstraintEngine._eval_bool_expr(
            "not x0", {"x0": False}
        ) is True

    def test_complex_expression(self):
        assert ConstraintEngine._eval_bool_expr(
            "not (x0 and x1 and x2)",
            {"x0": True, "x1": True, "x2": True},
        ) is False

    def test_variable_ordering(self):
        """Ensure x10 doesn't get partially replaced by x1."""
        result = ConstraintEngine._eval_bool_expr(
            "x1 and x10",
            {"x1": True, "x10": False},
        )
        assert result is False