""" Integration tests for Flexible Database Client with SQLite. """ import pytest import pytest_asyncio import tempfile import os from src.database.flexible_client import FlexibleDatabaseClient class TestFlexibleDatabaseClient: """Integration tests for flexible database client.""" @pytest_asyncio.fixture async def sqlite_db_client(self): """Create SQLite database client for testing.""" # Create temporary SQLite database temp_db = tempfile.NamedTemporaryFile(suffix='.db', delete=False) temp_db.close() client = FlexibleDatabaseClient(f"sqlite:///{temp_db.name}") # Connect and create tables await client.connect() client.create_tables() # Insert test data client.execute(""" INSERT INTO pump_stations (station_id, station_name, location) VALUES ('STATION_001', 'Main Pump Station', 'Downtown Area'), ('STATION_002', 'Secondary Station', 'Industrial Zone') """) client.execute(""" INSERT INTO pumps (station_id, pump_id, pump_name, control_type, min_speed_hz, max_speed_hz, default_setpoint_hz) VALUES ('STATION_001', 'PUMP_001', 'Main Pump 1', 'DIRECT_SPEED', 20.0, 60.0, 35.0), ('STATION_001', 'PUMP_002', 'Main Pump 2', 'LEVEL_CONTROLLED', 20.0, 60.0, 35.0), ('STATION_002', 'PUMP_001', 'Secondary Pump 1', 'POWER_CONTROLLED', 20.0, 60.0, 35.0) """) client.execute(""" INSERT INTO pump_plans ( station_id, pump_id, target_flow_m3h, target_power_kw, target_level_m, suggested_speed_hz, interval_start, interval_end, plan_version, plan_status, optimization_run_id ) VALUES ('STATION_001', 'PUMP_001', 150.0, NULL, NULL, 42.5, datetime('now', '-1 hour'), datetime('now', '+1 hour'), 1, 'ACTIVE', 'OPT_RUN_001'), ('STATION_001', 'PUMP_002', NULL, NULL, 2.5, 38.0, datetime('now', '-1 hour'), datetime('now', '+1 hour'), 1, 'ACTIVE', 'OPT_RUN_001'), ('STATION_002', 'PUMP_001', NULL, 18.5, NULL, 40.0, datetime('now', '-1 hour'), datetime('now', '+1 hour'), 1, 'ACTIVE', 'OPT_RUN_001') """) client.execute(""" INSERT INTO pump_feedback ( station_id, pump_id, actual_speed_hz, actual_power_kw, actual_flow_m3h, wet_well_level_m, pump_running, alarm_active ) VALUES ('STATION_001', 'PUMP_001', 42.5, 16.2, 148.5, 1.8, 1, 0), ('STATION_001', 'PUMP_002', 38.0, 14.8, 135.2, 2.3, 1, 0), ('STATION_002', 'PUMP_001', 40.0, 18.3, 142.1, 1.9, 1, 0) """) yield client # Clean up await client.disconnect() os.unlink(temp_db.name) def test_connect_sqlite(self, sqlite_db_client): """Test connecting to SQLite database.""" # Connection is established in fixture assert sqlite_db_client.health_check() is True stats = sqlite_db_client.get_connection_stats() assert stats["database_type"] == "SQLite" assert stats["status"] == "connected" def test_get_pump_stations(self, sqlite_db_client): """Test getting pump stations.""" stations = sqlite_db_client.get_pump_stations() assert len(stations) == 2 assert stations[0]["station_id"] == "STATION_001" assert stations[1]["station_id"] == "STATION_002" def test_get_pumps(self, sqlite_db_client): """Test getting pumps.""" # Get all pumps pumps = sqlite_db_client.get_pumps() assert len(pumps) == 3 # Get pumps for specific station station_pumps = sqlite_db_client.get_pumps("STATION_001") assert len(station_pumps) == 2 assert all(p["station_id"] == "STATION_001" for p in station_pumps) def test_get_pump(self, sqlite_db_client): """Test getting specific pump.""" pump = sqlite_db_client.get_pump("STATION_001", "PUMP_001") assert pump is not None assert pump["pump_id"] == "PUMP_001" assert pump["control_type"] == "DIRECT_SPEED" assert pump["min_speed_hz"] == 20.0 assert pump["max_speed_hz"] == 60.0 def test_get_current_plan(self, sqlite_db_client): """Test getting current plan.""" plan = sqlite_db_client.get_current_plan("STATION_001", "PUMP_001") assert plan is not None assert plan["suggested_speed_hz"] == 42.5 assert plan["plan_status"] == "ACTIVE" def test_get_latest_feedback(self, sqlite_db_client): """Test getting latest feedback.""" feedback = sqlite_db_client.get_latest_feedback("STATION_001", "PUMP_001") assert feedback is not None assert feedback["actual_speed_hz"] == 42.5 assert feedback["pump_running"] == 1 def test_get_pump_feedback(self, sqlite_db_client): """Test getting pump feedback.""" feedback = sqlite_db_client.get_pump_feedback("STATION_001", "PUMP_001", limit=2) assert len(feedback) == 1 # Only one record in test data assert feedback[0]["actual_speed_hz"] == 42.5 def test_execute_query(self, sqlite_db_client): """Test custom query execution.""" result = sqlite_db_client.execute_query( "SELECT COUNT(*) as count FROM pumps WHERE station_id = :station_id", {"station_id": "STATION_001"} ) assert result[0]["count"] == 2 def test_execute_update(self, sqlite_db_client): """Test update execution.""" rows_affected = sqlite_db_client.execute( "UPDATE pumps SET pump_name = :new_name WHERE station_id = :station_id AND pump_id = :pump_id", { "new_name": "Updated Pump Name", "station_id": "STATION_001", "pump_id": "PUMP_001" } ) assert rows_affected == 1 # Verify update pump = sqlite_db_client.get_pump("STATION_001", "PUMP_001") assert pump["pump_name"] == "Updated Pump Name" def test_health_check(self, sqlite_db_client): """Test health check.""" assert sqlite_db_client.health_check() is True def test_connection_stats(self, sqlite_db_client): """Test connection statistics.""" stats = sqlite_db_client.get_connection_stats() assert "database_type" in stats assert "pool_size" in stats assert "status" in stats assert stats["database_type"] == "SQLite" def test_error_handling(self, sqlite_db_client): """Test error handling.""" # Test with invalid query with pytest.raises(Exception): sqlite_db_client.execute_query("SELECT * FROM non_existent_table") # Test with non-existent pump pump = sqlite_db_client.get_pump("NON_EXISTENT", "PUMP_001") assert pump is None def test_create_tables_idempotent(self, sqlite_db_client): """Test that create_tables is idempotent.""" # Should not raise an exception when tables already exist sqlite_db_client.create_tables() # Verify tables still work stations = sqlite_db_client.get_pump_stations() assert len(stations) == 2