180 lines
7.2 KiB
Python
180 lines
7.2 KiB
Python
"""
|
|
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 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 |