test_z3.py•6.76 kB
from returns.result import Success
from usolver_mcp.models.z3_models import (
Z3Constraint,
Z3Problem,
Z3Variable,
Z3VariableType,
)
from usolver_mcp.solvers.z3_solver import solve_problem
pytest_plugins: list[str] = []
def test_worker_task_assignment_problem() -> None:
"""Test solving a worker-task assignment problem using Z3.
Problem Summary:
A simple scheduling problem where two workers (A and B) need to be assigned
to two tasks (1 and 2) with the following constraints:
- Each worker can only be assigned to one task
- Each task must be assigned to exactly one worker
- Worker A prefers task 1 (soft constraint/preference)
This demonstrates Z3's ability to handle boolean satisfiability problems
with logical constraints and preferences.
"""
# Define the problem
problem = Z3Problem(
variables=[
# Binary variables for worker assignments
Z3Variable(
name="A1", type=Z3VariableType.BOOLEAN
), # Worker A assigned to Task 1
Z3Variable(
name="A2", type=Z3VariableType.BOOLEAN
), # Worker A assigned to Task 2
Z3Variable(
name="B1", type=Z3VariableType.BOOLEAN
), # Worker B assigned to Task 1
Z3Variable(
name="B2", type=Z3VariableType.BOOLEAN
), # Worker B assigned to Task 2
],
constraints=[
# Worker A can only be assigned to one task
Z3Constraint(
expression="And(Or(A1, A2), Not(And(A1, A2)))",
description="Worker A must be assigned to exactly one task",
),
# Worker B can only be assigned to one task
Z3Constraint(
expression="And(Or(B1, B2), Not(And(B1, B2)))",
description="Worker B must be assigned to exactly one task",
),
# Task 1 must be assigned to one worker
Z3Constraint(
expression="And(Or(A1, B1), Not(And(A1, B1)))",
description="Task 1 must be assigned to exactly one worker",
),
# Task 2 must be assigned to one worker
Z3Constraint(
expression="And(Or(A2, B2), Not(And(A2, B2)))",
description="Task 2 must be assigned to exactly one worker",
),
# Preference: Worker A prefers task 1 (soft constraint)
Z3Constraint(
expression="A1",
description="Preference for Worker A to be assigned to Task 1",
),
],
description="Worker-Task Assignment Problem with Preferences",
)
# Solve the problem
result = solve_problem(problem)
# Verify the solution
assert isinstance(result, Success)
solution = result.unwrap()
# Check that the problem is satisfiable
assert solution.is_satisfiable
assert solution.status == "sat"
# Check that we have values for all 4 variables
assert len(solution.values) == 4
assert all(var in solution.values for var in ["A1", "A2", "B1", "B2"])
# Verify that all values are boolean
for var_name, value in solution.values.items():
assert isinstance(
value, bool
), f"Variable {var_name} should be boolean, got {type(value)}"
# Check constraint satisfaction
values = solution.values
# Worker A assigned to exactly one task
a_assignments = sum([values["A1"], values["A2"]])
assert a_assignments == 1, "Worker A must be assigned to exactly one task"
# Worker B assigned to exactly one task
b_assignments = sum([values["B1"], values["B2"]])
assert b_assignments == 1, "Worker B must be assigned to exactly one task"
# Task 1 assigned to exactly one worker
task1_assignments = sum([values["A1"], values["B1"]])
assert task1_assignments == 1, "Task 1 must be assigned to exactly one worker"
# Task 2 assigned to exactly one worker
task2_assignments = sum([values["A2"], values["B2"]])
assert task2_assignments == 1, "Task 2 must be assigned to exactly one worker"
# Check preference: Worker A should be assigned to Task 1 (due to preference constraint)
assert (
values["A1"] is True
), "Worker A should be assigned to Task 1 due to preference"
assert values["B2"] is True, "Worker B should be assigned to Task 2"
assert values["A2"] is False, "Worker A should not be assigned to Task 2"
assert values["B1"] is False, "Worker B should not be assigned to Task 1"
def test_simple_arithmetic_constraints() -> None:
"""Test Z3 with simple arithmetic constraints.
Problem: Find integer values for x and y such that:
- x + y = 10
- x - y = 2
- x > 0, y > 0
Expected solution: x = 6, y = 4
"""
problem = Z3Problem(
variables=[
Z3Variable(name="x", type=Z3VariableType.INTEGER),
Z3Variable(name="y", type=Z3VariableType.INTEGER),
],
constraints=[
Z3Constraint(expression="x + y == 10", description="Sum equals 10"),
Z3Constraint(expression="x - y == 2", description="Difference equals 2"),
Z3Constraint(expression="x > 0", description="x is positive"),
Z3Constraint(expression="y > 0", description="y is positive"),
],
description="Simple arithmetic constraint problem",
)
result = solve_problem(problem)
assert isinstance(result, Success)
solution = result.unwrap()
assert solution.is_satisfiable
# Check the expected solution
assert solution.values["x"] == 6
assert solution.values["y"] == 4
# Verify constraints are satisfied
x, y = solution.values["x"], solution.values["y"]
assert x + y == 10
assert x - y == 2
assert x > 0
assert y > 0
def test_unsatisfiable_problem() -> None:
"""Test Z3 with an unsatisfiable problem.
Problem: Find a boolean variable x such that:
- x is true
- x is false
This should be unsatisfiable.
"""
problem = Z3Problem(
variables=[
Z3Variable(name="x", type=Z3VariableType.BOOLEAN),
],
constraints=[
Z3Constraint(expression="x", description="x must be true"),
Z3Constraint(expression="Not(x)", description="x must be false"),
],
description="Unsatisfiable boolean problem",
)
result = solve_problem(problem)
assert isinstance(result, Success)
solution = result.unwrap()
# Should be unsatisfiable
assert not solution.is_satisfiable
assert solution.status == "unsat"
assert len(solution.values) == 0 # No solution values for unsatisfiable problems