CalejoControl/tests/unit/test_setpoint_manager.py

"""
Unit tests for SetpointManager and calculators.
"""

import pytest
from unittest.mock import Mock, patch
from datetime import datetime

from src.core.setpoint_manager import (
    SetpointManager, 
    DirectSpeedCalculator, 
    LevelControlledCalculator, 
    PowerControlledCalculator
)


class TestSetpointCalculators:
    """Test cases for setpoint calculators."""
    
    def test_direct_speed_calculator(self):
        """Test direct speed calculator."""
        calculator = DirectSpeedCalculator()
        
        # Test with suggested speed
        plan = {'suggested_speed_hz': 42.5}
        feedback = None
        pump_info = {}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 42.5
        
        # Test without suggested speed (fallback)
        plan = {}
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 35.0
    
    def test_level_controlled_calculator_with_feedback(self):
        """Test level controlled calculator with feedback."""
        calculator = LevelControlledCalculator()
        
        # Test with level feedback (target > current)
        plan = {'target_level_m': 3.0, 'suggested_speed_hz': 40.0}
        feedback = {'current_level_m': 2.0}
        pump_info = {}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        # Expected: 35.0 + 5.0 * (3.0 - 2.0) = 40.0
        assert result == 40.0
        
        # Test with level feedback (target < current)
        plan = {'target_level_m': 2.0, 'suggested_speed_hz': 40.0}
        feedback = {'current_level_m': 3.0}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        # Expected: 35.0 + 5.0 * (2.0 - 3.0) = 30.0
        assert result == 30.0
    
    def test_level_controlled_calculator_without_feedback(self):
        """Test level controlled calculator without feedback."""
        calculator = LevelControlledCalculator()
        
        # Test without feedback (fallback to suggested speed)
        plan = {'suggested_speed_hz': 38.5}
        feedback = None
        pump_info = {}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 38.5
        
        # Test without suggested speed (fallback to default)
        plan = {}
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 35.0
    
    def test_power_controlled_calculator_with_feedback(self):
        """Test power controlled calculator with feedback."""
        calculator = PowerControlledCalculator()
        
        # Test with power feedback (target > current)
        plan = {'target_power_kw': 20.0, 'suggested_speed_hz': 40.0}
        feedback = {'current_power_kw': 15.0}
        pump_info = {}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        # Expected: 35.0 + 2.0 * (20.0 - 15.0) = 45.0
        assert result == 45.0
        
        # Test with power feedback (target < current)
        plan = {'target_power_kw': 15.0, 'suggested_speed_hz': 40.0}
        feedback = {'current_power_kw': 20.0}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        # Expected: 35.0 + 2.0 * (15.0 - 20.0) = 25.0
        assert result == 25.0
    
    def test_power_controlled_calculator_without_feedback(self):
        """Test power controlled calculator without feedback."""
        calculator = PowerControlledCalculator()
        
        # Test without feedback (fallback to suggested speed)
        plan = {'suggested_speed_hz': 37.5}
        feedback = None
        pump_info = {}
        
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 37.5
        
        # Test without suggested speed (fallback to default)
        plan = {}
        result = calculator.calculate_setpoint(plan, feedback, pump_info)
        assert result == 35.0


class TestSetpointManager:
    """Test cases for SetpointManager."""
    
    def setup_method(self):
        """Set up test fixtures."""
        self.mock_discovery = Mock()
        self.mock_db_client = Mock()
        self.mock_safety_enforcer = Mock()
        self.mock_emergency_stop_manager = Mock()
        self.mock_watchdog = Mock()
        
        # Configure mocks
        self.mock_safety_enforcer.enforce_limits = Mock(return_value=40.0)
        self.mock_emergency_stop_manager.is_emergency_stop_active = Mock(return_value=False)
        self.mock_watchdog.is_failsafe_active = Mock(return_value=False)
        
        self.setpoint_manager = SetpointManager(
            discovery=self.mock_discovery,
            db_client=self.mock_db_client,
            safety_enforcer=self.mock_safety_enforcer,
            emergency_stop_manager=self.mock_emergency_stop_manager,
            watchdog=self.mock_watchdog
        )
    
    def test_get_current_setpoint_normal_operation(self):
        """Test setpoint calculation in normal operation."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        pump_info = {
            'station_id': station_id,
            'pump_id': pump_id,
            'control_type': 'DIRECT_SPEED'
        }
        
        plan = {
            'suggested_speed_hz': 42.5
        }
        
        self.mock_discovery.get_pump.return_value = pump_info
        self.mock_db_client.get_current_plan.return_value = plan
        self.mock_db_client.get_latest_feedback.return_value = None
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 40.0  # After safety enforcement
        self.mock_discovery.get_pump.assert_called_once_with(station_id, pump_id)
        self.mock_db_client.get_current_plan.assert_called_once_with(station_id, pump_id)
        self.mock_safety_enforcer.enforce_limits.assert_called_once_with(station_id, pump_id, 42.5)
    
    def test_get_current_setpoint_emergency_stop(self):
        """Test setpoint calculation during emergency stop."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_emergency_stop_manager.is_emergency_stop_active.return_value = True
        self.mock_db_client.execute_query.return_value = [{'default_setpoint_hz': 30.0}]
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 30.0
        self.mock_emergency_stop_manager.is_emergency_stop_active.assert_called_once_with(station_id, pump_id)
        # Should not call other methods during emergency stop
        self.mock_discovery.get_pump.assert_not_called()
        self.mock_db_client.get_current_plan.assert_not_called()
    
    def test_get_current_setpoint_failsafe_mode(self):
        """Test setpoint calculation during failsafe mode."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_watchdog.is_failsafe_active.return_value = True
        self.mock_db_client.execute_query.return_value = [{'default_setpoint_hz': 25.0}]
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 25.0
        self.mock_watchdog.is_failsafe_active.assert_called_once_with(station_id, pump_id)
        # Should not call other methods during failsafe mode
        self.mock_discovery.get_pump.assert_not_called()
        self.mock_db_client.get_current_plan.assert_not_called()
    
    def test_get_current_setpoint_no_pump_found(self):
        """Test setpoint calculation when pump is not found."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_discovery.get_pump.return_value = None
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result is None
        self.mock_discovery.get_pump.assert_called_once_with(station_id, pump_id)
    
    def test_get_current_setpoint_no_active_plan(self):
        """Test setpoint calculation when no active plan exists."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        pump_info = {
            'station_id': station_id,
            'pump_id': pump_id,
            'control_type': 'DIRECT_SPEED'
        }
        
        self.mock_discovery.get_pump.return_value = pump_info
        self.mock_db_client.get_current_plan.return_value = None
        self.mock_db_client.execute_query.return_value = [{'default_setpoint_hz': 35.0}]
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 35.0  # Default setpoint
        self.mock_discovery.get_pump.assert_called_once_with(station_id, pump_id)
        self.mock_db_client.get_current_plan.assert_called_once_with(station_id, pump_id)
    
    def test_get_current_setpoint_unknown_control_type(self):
        """Test setpoint calculation with unknown control type."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        pump_info = {
            'station_id': station_id,
            'pump_id': pump_id,
            'control_type': 'UNKNOWN_TYPE'
        }
        
        self.mock_discovery.get_pump.return_value = pump_info
        self.mock_db_client.get_current_plan.return_value = {'suggested_speed_hz': 40.0}
        
        # Act
        result = self.setpoint_manager.get_current_setpoint(station_id, pump_id)
        
        # Assert
        assert result is None
        self.mock_discovery.get_pump.assert_called_once_with(station_id, pump_id)
        self.mock_db_client.get_current_plan.assert_called_once_with(station_id, pump_id)
    
    def test_get_all_current_setpoints(self):
        """Test getting setpoints for all pumps."""
        # Arrange
        stations = [
            {'station_id': 'STATION_001'},
            {'station_id': 'STATION_002'}
        ]
        
        pumps_station_001 = [
            {'pump_id': 'PUMP_001'},
            {'pump_id': 'PUMP_002'}
        ]
        
        pumps_station_002 = [
            {'pump_id': 'PUMP_001'}
        ]
        
        self.mock_discovery.get_stations.return_value = stations
        self.mock_discovery.get_pumps.side_effect = [pumps_station_001, pumps_station_002]
        
        # Mock get_current_setpoint to return different values
        def mock_get_current_setpoint(station_id, pump_id):
            return float(f"{ord(station_id[-1])}.{ord(pump_id[-1])}")
        
        self.setpoint_manager.get_current_setpoint = Mock(side_effect=mock_get_current_setpoint)
        
        # Act
        result = self.setpoint_manager.get_all_current_setpoints()
        
        # Assert
        assert 'STATION_001' in result
        assert 'STATION_002' in result
        assert 'PUMP_001' in result['STATION_001']
        assert 'PUMP_002' in result['STATION_001']
        assert 'PUMP_001' in result['STATION_002']
        
        # Verify all pumps were queried
        assert self.setpoint_manager.get_current_setpoint.call_count == 3
    
    def test_get_default_setpoint_from_database(self):
        """Test getting default setpoint from database."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_db_client.execute_query.return_value = [{'default_setpoint_hz': 32.5}]
        
        # Act
        result = self.setpoint_manager._get_default_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 32.5
        self.mock_db_client.execute_query.assert_called_once()
    
    def test_get_default_setpoint_fallback(self):
        """Test getting default setpoint fallback when database fails."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_db_client.execute_query.return_value = []
        
        # Act
        result = self.setpoint_manager._get_default_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 35.0  # Conservative fallback
        self.mock_db_client.execute_query.assert_called_once()
    
    def test_get_default_setpoint_database_error(self):
        """Test getting default setpoint when database query fails."""
        # Arrange
        station_id = 'STATION_001'
        pump_id = 'PUMP_001'
        
        self.mock_db_client.execute_query.side_effect = Exception("Database error")
        
        # Act
        result = self.setpoint_manager._get_default_setpoint(station_id, pump_id)
        
        # Assert
        assert result == 35.0  # Conservative fallback
        self.mock_db_client.execute_query.assert_called_once()