d7cd9043b2
Add unique constraint on (model_id, timestamp, variable) in schema and implement idempotent export test to ensure no duplicates are created. Update CoreScouter and related tests to handle conflict resolution by ignoring duplicates.
364 lines
13 KiB
Python
364 lines
13 KiB
Python
"""
|
|
Pytest configuration and fixtures for E2E tests.
|
|
"""
|
|
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import pytest_asyncio
|
|
from pandas import DataFrame
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.orm import sessionmaker
|
|
from testcontainers.postgres import PostgresContainer
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from scouter.activities.activities import Activities
|
|
from scouter.workflow.pi_web_api_scouter import PIWebAPIScouter
|
|
from scouter.workflow.sub_workflows.core_scouter import CoreScouter
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
|
|
from e2e.fixtures.fake_mongodb_repository import FakeMongoDBRepository
|
|
from e2e.fixtures.fake_redis_repository import FakeRedisRepository
|
|
|
|
# Test constants
|
|
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
|
|
TEST_DATABASE_NAME = 'test_db'
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='session')
|
|
def postgres_container():
|
|
"""
|
|
Create a PostgreSQL container using testcontainers.
|
|
|
|
This fixture creates a real PostgreSQL database in a Docker container
|
|
that will be used for all tests in the session.
|
|
"""
|
|
postgres = PostgresContainer('postgres:15')
|
|
postgres.start()
|
|
yield postgres
|
|
postgres.stop()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def postgres_engine(postgres_container):
|
|
"""
|
|
Create SQLAlchemy engine for PostgreSQL test database.
|
|
|
|
This fixture creates a connection to the PostgreSQL container
|
|
created by the postgres_container fixture.
|
|
"""
|
|
# Get connection URL from container
|
|
connection_string = postgres_container.get_connection_url()
|
|
engine = create_engine(connection_string)
|
|
|
|
yield engine
|
|
|
|
engine.dispose()
|
|
|
|
|
|
def _create_schema_and_table(engine):
|
|
"""
|
|
Helper function to create schema and table in the given engine.
|
|
|
|
This is used by both the autouse fixture and test_activities to ensure
|
|
the schema exists before Activities tries to use it.
|
|
|
|
Note: For tests, we create a non-partitioned table to avoid issues
|
|
with pandas to_sql recognizing partitioned tables.
|
|
"""
|
|
schema_name = 'sientia_data'
|
|
table_name = 'laborious_data'
|
|
|
|
# Use begin() to ensure transaction is properly committed
|
|
with engine.begin() as conn:
|
|
# Create schema
|
|
conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema_name}"))
|
|
|
|
# Create table WITHOUT partitioning (simpler for tests)
|
|
# Same structure as production, but without PARTITION BY RANGE
|
|
create_table_sql = f"""
|
|
CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (
|
|
id SERIAL NOT NULL,
|
|
model_id int4 NOT NULL,
|
|
variable text NOT NULL,
|
|
value numeric NULL,
|
|
"timestamp" timestamptz NOT NULL,
|
|
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
|
CONSTRAINT unique_timestamp_variable UNIQUE (model_id, "timestamp", variable),
|
|
PRIMARY KEY (id, created_at)
|
|
);
|
|
"""
|
|
|
|
conn.execute(text(create_table_sql))
|
|
# Transaction is automatically committed when exiting the 'with' block
|
|
|
|
|
|
@pytest_asyncio.fixture(autouse=True)
|
|
def setup_postgres_schema_and_table(postgres_engine):
|
|
"""
|
|
Automatically create necessary schema and table before each test.
|
|
|
|
This fixture runs automatically (autouse=True) and ensures
|
|
that the sientia_data schema and laborious_data table exist
|
|
with the correct structure before tests execute.
|
|
|
|
Note: For tests, we use a non-partitioned table which is simpler and avoids issues
|
|
with pandas to_sql recognizing partitioned tables.
|
|
"""
|
|
_create_schema_and_table(postgres_engine)
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_logger():
|
|
"""Mock logger for testing."""
|
|
logger = MagicMock(spec=Logger)
|
|
logger.info = MagicMock()
|
|
logger.debug = MagicMock()
|
|
logger.error = MagicMock()
|
|
logger.warning = MagicMock()
|
|
logger.custom_info = MagicMock()
|
|
return logger
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_mongo_client():
|
|
"""
|
|
Mock MongoDB client to avoid real connections.
|
|
|
|
This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
|
|
allowing us to use a real NotificationHandler instance without connecting to MongoDB.
|
|
"""
|
|
mock_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
|
|
# Configure the mock chain: client[database] -> db[collection] -> collection
|
|
mock_client.__getitem__.return_value = mock_db
|
|
mock_db.__getitem__.return_value = mock_collection
|
|
|
|
# Mock server_info() to avoid connection attempts
|
|
mock_client.server_info = MagicMock()
|
|
|
|
# Mock insert_one for notifications
|
|
mock_collection.insert_one = MagicMock()
|
|
|
|
return mock_client
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def notification_handler(mock_logger, mock_mongo_client):
|
|
"""
|
|
Create a real NotificationHandler instance with mocked MongoDB client.
|
|
|
|
This fixture creates a real CoreNotificationHandler instance but mocks
|
|
the underlying MongoDB connection to avoid real database connections.
|
|
"""
|
|
# Patch MongoClient where it's imported in the handlers module
|
|
with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
|
|
handler = CoreNotificationHandler(
|
|
connection_string=TEST_MONGODB_CONNECTION_STRING,
|
|
database=TEST_DATABASE_NAME,
|
|
logger=mock_logger,
|
|
project_name='scouter',
|
|
)
|
|
yield handler
|
|
handler.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def metrics_controller(mock_logger):
|
|
"""
|
|
Create a real MetricsController instance.
|
|
|
|
MetricsController doesn't require external services, so we can use
|
|
a real instance without mocking anything.
|
|
"""
|
|
controller = MetricsController(logger=mock_logger)
|
|
yield controller
|
|
# MetricsController might have cleanup, but it's optional
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_pi_web_api_client():
|
|
"""Mock PI Web API client."""
|
|
mock_client = MagicMock()
|
|
|
|
# Mock DataFrame response similar to real API
|
|
# The real PIWebAPIClient returns timestamp as datetime, so we need to match that
|
|
mock_df = DataFrame({
|
|
'timestamp': [
|
|
'2024-01-01 12:00:00+0000',
|
|
'2024-01-01 12:01:00+0000',
|
|
'2024-01-01 12:02:00+0000',
|
|
],
|
|
'name': ['tag1', 'tag2', 'tag3'],
|
|
'value': [10.5, 20.3, 30.7],
|
|
'tag': ['webid1', 'webid2', 'webid3'],
|
|
})
|
|
|
|
# Convert timestamp to datetime (UTC, floored to seconds) to match real client behavior
|
|
mock_df['timestamp'] = pd.to_datetime(mock_df['timestamp'], utc=True).dt.floor('s')
|
|
|
|
mock_client.get_latest_values_df = AsyncMock(return_value=mock_df)
|
|
mock_client.close = MagicMock()
|
|
return mock_client
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
def test_activities(
|
|
postgres_engine,
|
|
postgres_container,
|
|
mock_logger,
|
|
notification_handler,
|
|
metrics_controller,
|
|
mock_pi_web_api_client,
|
|
):
|
|
"""
|
|
Create Activities instance with test dependencies.
|
|
|
|
This fixture creates a real Activities instance with:
|
|
- PostgreSQL database (via testcontainers)
|
|
- FakeRedis instead of real Redis
|
|
- FakeMongoDB instead of real MongoDB
|
|
- Mocked PI Web API client
|
|
- Real NotificationHandler and MetricsController (with mocked underlying services)
|
|
"""
|
|
# Get connection details from container
|
|
connection_string = postgres_container.get_connection_url()
|
|
|
|
# Parse connection string to get individual components
|
|
# Format: postgresql://testuser:testpass@localhost:5432/test
|
|
from urllib.parse import urlparse
|
|
parsed = urlparse(connection_string)
|
|
|
|
# Ensure schema and table exist BEFORE creating Activities
|
|
# This ensures the schema exists when Activities initializes its engine
|
|
_create_schema_and_table(postgres_engine)
|
|
|
|
# Create fake repositories that will be used from the start
|
|
fake_redis_repo = FakeRedisRepository(
|
|
host='localhost',
|
|
port=6379,
|
|
username='',
|
|
password='',
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
fake_mongo_repo = FakeMongoDBRepository(
|
|
connection_string=TEST_MONGODB_CONNECTION_STRING,
|
|
database_name=TEST_DATABASE_NAME,
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
# Patch create_engine to return our postgres_engine instead of creating a new one
|
|
# This ensures Activities uses the same engine from the start
|
|
original_create_engine = create_engine
|
|
|
|
def patched_create_engine(connection_string, *args, **kwargs):
|
|
# Check if this is the connection string that Activities would create
|
|
# Activities creates: postgresql://user:password@host:port/dbname
|
|
expected_conn_str = f"postgresql://{parsed.username or 'test'}:{parsed.password or 'test'}@localhost:{postgres_container.get_exposed_port(5432)}/{parsed.path.lstrip('/') if parsed.path else 'test'}"
|
|
|
|
# If it matches our test container connection, return our engine
|
|
if connection_string == expected_conn_str:
|
|
return postgres_engine
|
|
# Otherwise, use the original create_engine
|
|
return original_create_engine(connection_string, *args, **kwargs)
|
|
|
|
# Patch RedisRepository to return our fake repository from the start
|
|
def patched_redis_repository(*args, **kwargs):
|
|
return fake_redis_repo
|
|
|
|
# Patch MongoDBRepository to return our fake repository from the start
|
|
def patched_mongodb_repository(*args, **kwargs):
|
|
return fake_mongo_repo
|
|
|
|
# Patch PIWebAPIClient to return our mock from the start
|
|
def patched_pi_web_api_client(*args, **kwargs):
|
|
return mock_pi_web_api_client
|
|
|
|
# Create Activities with test configurations
|
|
# All patches ensure it uses our test instances from the start
|
|
with patch('sientia_do.temporal.activities.postgres.create_engine', new=patched_create_engine), \
|
|
patch('scouter.activities.redis.RedisRepository', new=patched_redis_repository), \
|
|
patch('scouter.activities.mongodb.MongoDBRepository', new=patched_mongodb_repository), \
|
|
patch('scouter.activities.api.PIWebAPIClient', new=patched_pi_web_api_client):
|
|
|
|
activities = Activities(
|
|
postgres_config={
|
|
'host': 'localhost', # Container exposes to localhost
|
|
'port': postgres_container.get_exposed_port(5432),
|
|
'user': parsed.username or 'test',
|
|
'password': parsed.password or 'test',
|
|
'dbname': parsed.path.lstrip('/') if parsed.path else 'test',
|
|
'min_connections': 1,
|
|
'max_connections': 5,
|
|
},
|
|
redis_config={
|
|
'host': 'localhost',
|
|
'port': 6379,
|
|
'username': '',
|
|
'password': '',
|
|
},
|
|
mongodb_config={
|
|
'connection_string': TEST_MONGODB_CONNECTION_STRING,
|
|
'database_name': TEST_DATABASE_NAME,
|
|
},
|
|
api_config={
|
|
'base_url': 'http://localhost:8080',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
},
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
# Verify that Activities is using our instances from the start
|
|
assert activities.engine is postgres_engine, "Activities should use the same engine as postgres_engine"
|
|
assert activities.redis_repository is fake_redis_repo, "Activities should use the same fake Redis repository"
|
|
assert activities.mongodb_repository is fake_mongo_repo, "Activities should use the same fake MongoDB repository"
|
|
assert activities.pi_web_api_client is mock_pi_web_api_client, "Activities should use the same mock PI Web API client"
|
|
|
|
yield activities
|
|
|
|
# Cleanup
|
|
activities.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def temporal_test_env():
|
|
"""Create Temporal test environment."""
|
|
env = await WorkflowEnvironment.start_time_skipping()
|
|
async with env:
|
|
yield env
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def temporal_worker(temporal_test_env, test_activities):
|
|
"""Create Temporal worker with test activities."""
|
|
async with Worker(
|
|
temporal_test_env.client,
|
|
task_queue='test-queue',
|
|
workflows=[PIWebAPIScouter, CoreScouter],
|
|
activities=[
|
|
test_activities.get_tag_values,
|
|
test_activities.data_quality_gate,
|
|
test_activities.aggregate_data,
|
|
test_activities.group_and_hold_data,
|
|
test_activities.export_data_to_postgres,
|
|
test_activities.write_metrics,
|
|
test_activities.store_data_package,
|
|
],
|
|
) as worker:
|
|
yield worker
|