""" 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()