brain/tests/test_ml_effort_agent.py

"""
Unit tests for MLEffortAgent.
Tests effort matrix computation and breakdown calculation.
"""

import pytest
from unittest.mock import MagicMock
from uuid import uuid4

from app.services.ml_effort_agent import MLEffortAgent
from app.schemas.agent_schemas import EffortWeights


class MockDriver:
    """Mock Driver for testing."""
    def __init__(self, id=None, vehicle_capacity_kg=100.0, is_ev=False, battery_range_km=None, charging_time_minutes=0):
        self.id = id or uuid4()
        self.vehicle_capacity_kg = vehicle_capacity_kg
        self.is_ev = is_ev
        self.vehicle_type = "EV" if is_ev else "ICE"
        self.battery_range_km = battery_range_km
        self.charging_time_minutes = charging_time_minutes


class MockRoute:
    """Mock Route for testing."""
    def __init__(
        self,
        id=None,
        num_packages=10,
        total_weight_kg=20.0,
        num_stops=5,
        route_difficulty_score=2.0,
        estimated_time_minutes=60,
        total_distance_km=25.0,
        charging_time_minutes=0,
    ):
        self.id = id or uuid4()
        self.num_packages = num_packages
        self.total_weight_kg = total_weight_kg
        self.num_stops = num_stops
        self.route_difficulty_score = route_difficulty_score
        self.estimated_time_minutes = estimated_time_minutes
        self.total_distance_km = total_distance_km
        self.charging_time_minutes = charging_time_minutes


class TestMLEffortAgent:
    """Test suite for MLEffortAgent."""
    
    def test_effort_matrix_shape(self):
        """Test that effort matrix has correct dimensions."""
        agent = MLEffortAgent()
        
        drivers = [MockDriver() for _ in range(3)]
        routes = [MockRoute() for _ in range(4)]
        
        result = agent.compute_effort_matrix(drivers, routes)
        
        assert len(result.matrix) == 3, "Should have 3 rows (drivers)"
        assert all(len(row) == 4 for row in result.matrix), "Each row should have 4 columns (routes)"
        assert len(result.driver_ids) == 3
        assert len(result.route_ids) == 4
    
    def test_effort_increases_with_packages(self):
        """Test that effort increases with more packages."""
        agent = MLEffortAgent()
        
        driver = MockDriver()
        route_light = MockRoute(num_packages=5, total_weight_kg=10.0)
        route_heavy = MockRoute(num_packages=20, total_weight_kg=40.0)
        
        result_light = agent.compute_effort_matrix([driver], [route_light])
        result_heavy = agent.compute_effort_matrix([driver], [route_heavy])
        
        effort_light = result_light.matrix[0][0]
        effort_heavy = result_heavy.matrix[0][0]
        
        assert effort_heavy > effort_light, "Heavier route should have more effort"
    
    def test_effort_increases_with_difficulty(self):
        """Test that effort increases with route difficulty."""
        agent = MLEffortAgent()
        
        driver = MockDriver()
        route_easy = MockRoute(route_difficulty_score=1.0)
        route_hard = MockRoute(route_difficulty_score=5.0)
        
        result_easy = agent.compute_effort_matrix([driver], [route_easy])
        result_hard = agent.compute_effort_matrix([driver], [route_hard])
        
        assert result_hard.matrix[0][0] > result_easy.matrix[0][0]
    
    def test_capacity_penalty_applied(self):
        """Test that overloaded routes get penalty."""
        agent = MLEffortAgent()
        
        # Driver with 50kg capacity
        driver = MockDriver(vehicle_capacity_kg=50.0)
        
        # Routes with different weights
        route_under = MockRoute(total_weight_kg=40.0)  # Under capacity
        route_over = MockRoute(total_weight_kg=70.0)   # Over capacity
        
        result_under = agent.compute_effort_matrix([driver], [route_under])
        result_over = agent.compute_effort_matrix([driver], [route_over])
        
        # Over-capacity should have significantly higher effort
        assert result_over.matrix[0][0] > result_under.matrix[0][0]
    
    def test_breakdown_components(self):
        """Test that breakdown contains all components."""
        agent = MLEffortAgent()
        
        driver = MockDriver()
        route = MockRoute()
        
        result = agent.compute_effort_matrix([driver], [route])
        
        key = f"{driver.id}:{route.id}"
        assert key in result.breakdown
        
        breakdown = result.breakdown[key]
        assert hasattr(breakdown, 'physical_effort')
        assert hasattr(breakdown, 'route_complexity')
        assert hasattr(breakdown, 'time_pressure')
        assert hasattr(breakdown, 'capacity_penalty')
        assert hasattr(breakdown, 'total')
        
        # Total should be sum of components
        expected_total = (
            breakdown.physical_effort +
            breakdown.route_complexity +
            breakdown.time_pressure +
            breakdown.capacity_penalty
        )
        assert abs(breakdown.total - expected_total) < 0.01
    
    def test_stats_computed(self):
        """Test that matrix stats are computed correctly."""
        agent = MLEffortAgent()
        
        drivers = [MockDriver() for _ in range(2)]
        routes = [MockRoute(num_packages=i*5 + 5) for i in range(3)]
        
        result = agent.compute_effort_matrix(drivers, routes)
        
        assert "min" in result.stats
        assert "max" in result.stats
        assert "avg" in result.stats
        assert result.stats["min"] <= result.stats["avg"] <= result.stats["max"]
    
    def test_custom_weights(self):
        """Test that custom weights affect effort calculation."""
        # Default weights
        agent_default = MLEffortAgent()
        
        # Custom weights with higher package weight
        custom_weights = EffortWeights(alpha_packages=5.0, beta_weight=0.1)
        agent_custom = MLEffortAgent(weights=custom_weights)
        
        driver = MockDriver()
        route = MockRoute(num_packages=20, total_weight_kg=10.0)
        
        result_default = agent_default.compute_effort_matrix([driver], [route])
        result_custom = agent_custom.compute_effort_matrix([driver], [route])
        
        # Custom should be different due to different weights
        assert result_default.matrix[0][0] != result_custom.matrix[0][0]
    
    def test_empty_inputs(self):
        """Test handling of empty inputs."""
        agent = MLEffortAgent()
        
        result = agent.compute_effort_matrix([], [])
        
        assert result.matrix == []
        assert result.stats["min"] == 0.0
        assert result.stats["max"] == 0.0
        assert result.stats["avg"] == 0.0
    
    def test_snapshot_generation(self):
        """Test input/output snapshot generation for logging."""
        agent = MLEffortAgent()
        
        drivers = [MockDriver() for _ in range(2)]
        routes = [MockRoute() for _ in range(3)]
        
        input_snapshot = agent.get_input_snapshot(drivers, routes)
        assert input_snapshot["num_drivers"] == 2
        assert input_snapshot["num_routes"] == 3
        
        result = agent.compute_effort_matrix(drivers, routes)
        output_snapshot = agent.get_output_snapshot(result)
        assert "matrix_shape" in output_snapshot
        assert "min_effort" in output_snapshot