""" 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.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.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 # Use UNIQUE constraint directly since table is not partitioned 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, PRIMARY KEY (id, created_at), UNIQUE (model_id, timestamp, variable) ); """ conn.execute(text(create_table_sql)) # Transaction is automatically committed when exiting the 'with' block @pytest.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 with a UNIQUE constraint directly in the table definition, which is simpler and avoids issues with pandas to_sql recognizing partitioned tables. """ _create_schema_and_table(postgres_engine) yield @pytest.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.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.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.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.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 async 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 async def temporal_test_env(): """Create Temporal test environment.""" env = await WorkflowEnvironment.start_time_skipping() async with env: yield env @pytest_asyncio.fixture 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