vitor-aignosi
c12f2996c2
QTZPOC-13: Introduce joblib fallback for MLflow model loading Enhances model loading resilience by attempting to deserialize models directly with joblib if MLflow's native flavor loaders fail. This mechanism handles cases where .pkl files are present in the artifact. Adds joblib as a core dependency and includes new end-to-end and unit tests to validate fallback behavior, covering success, no-fallback, and double-failure scenarios. Also updates .gitignore for .claude/ files and incorporates minor refactorings and linting adjustments across E2E tests and utility functions.
768 lines
25 KiB
Python
768 lines
25 KiB
Python
"""
|
|
Pytest configuration and fixtures for E2E tests.
|
|
"""
|
|
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
|
|
_model_analysis_module = MagicMock()
|
|
_model_analysis_module.ModelAnalysis = MagicMock
|
|
sys.modules.setdefault('sientia', MagicMock())
|
|
sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
|
|
from io import BytesIO
|
|
|
|
import joblib
|
|
import pandas as pd
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sklearn.compose import ColumnTransformer
|
|
from sklearn.linear_model import LinearRegression
|
|
from sklearn.pipeline import Pipeline
|
|
from sqlalchemy import create_engine, text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
from testcontainers.minio import MinioContainer
|
|
from testcontainers.postgres import PostgresContainer
|
|
|
|
from e2e.opc_test_server import OpcE2ETestServer
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
|
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
|
|
|
# Test constants
|
|
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
|
|
TEST_DATABASE_NAME = 'test_db'
|
|
|
|
# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
|
|
JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
|
|
JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='session')
|
|
def minio_container():
|
|
"""
|
|
MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
|
|
"""
|
|
minio = MinioContainer()
|
|
minio.start()
|
|
yield minio
|
|
minio.stop()
|
|
|
|
|
|
@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.
|
|
"""
|
|
engine = create_engine(postgres_container.get_connection_url())
|
|
|
|
yield engine
|
|
|
|
engine.dispose()
|
|
|
|
|
|
def _create_schema_and_tables(engine):
|
|
"""
|
|
Helper function to create schema and tables in the given engine.
|
|
|
|
Creates predictions_schema with:
|
|
- laborious_data: Input data table for queries
|
|
- predictions: Output predictions table
|
|
- transformed_data: Output transformed data table
|
|
"""
|
|
# Use begin() to ensure transaction is properly committed
|
|
with engine.begin() as conn:
|
|
# Create predictions_schema
|
|
conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))
|
|
|
|
# Create laborious_data table (input data from sensors)
|
|
create_laborious_data_sql = """
|
|
CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
|
|
id SERIAL NOT NULL,
|
|
model_id int4 NOT NULL,
|
|
variable text NOT NULL,
|
|
value numeric NULL,
|
|
"timestamp" timestamptz NOT NULL,
|
|
created_at timestamptz NOT NULL,
|
|
PRIMARY KEY (id)
|
|
);
|
|
"""
|
|
conn.execute(text(create_laborious_data_sql))
|
|
|
|
# Create predictions table
|
|
create_predictions_sql = """
|
|
CREATE TABLE if not exists predictions_schema.predictions (
|
|
id SERIAL NOT NULL ,
|
|
model_id int4 NOT NULL,
|
|
prediction numeric NULL,
|
|
prediction_confidence numeric NOT NULL,
|
|
response_time numeric NOT NULL,
|
|
prediction_status text NOT NULL,
|
|
"timestamp" timestamptz NOT NULL,
|
|
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
|
"comments" text NULL,
|
|
PRIMARY KEY (id, created_at)
|
|
);
|
|
"""
|
|
conn.execute(text(create_predictions_sql))
|
|
|
|
# Create transformed_data table
|
|
create_transformed_sql = """
|
|
CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
|
|
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)
|
|
);
|
|
"""
|
|
conn.execute(text(create_transformed_sql))
|
|
|
|
|
|
@pytest_asyncio.fixture(autouse=True)
|
|
def setup_postgres_schema_and_tables(postgres_engine):
|
|
"""
|
|
Automatically create necessary schema and tables before each test.
|
|
|
|
This fixture runs automatically (autouse=True) and ensures
|
|
that the predictions_schema and tables exist with the correct structure.
|
|
"""
|
|
_create_schema_and_tables(postgres_engine)
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_logger():
|
|
"""Mock logger for testing."""
|
|
|
|
def message(message):
|
|
print(f'[LOG] {message}')
|
|
|
|
def custom_message(message, _metadata=None):
|
|
print(f'[LOG] {message}')
|
|
|
|
logger = MagicMock()
|
|
logger.info = MagicMock(side_effect=message)
|
|
logger.debug = MagicMock(side_effect=message)
|
|
logger.error = MagicMock(side_effect=message)
|
|
logger.warning = MagicMock(side_effect=message)
|
|
logger.custom_info = MagicMock(side_effect=custom_message)
|
|
logger.custom_debug = MagicMock(side_effect=custom_message)
|
|
logger.custom_error = MagicMock(side_effect=custom_message)
|
|
logger.custom_warning = MagicMock(side_effect=custom_message)
|
|
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.fixture
|
|
def notification_inserts(mock_mongo_client):
|
|
"""
|
|
Mongo insert_one mock used by CoreNotificationHandler for notification persistence.
|
|
|
|
Yields:
|
|
MagicMock for insert_one, reset before each test.
|
|
"""
|
|
mock_db = mock_mongo_client.__getitem__.return_value
|
|
mock_collection = mock_db.__getitem__.return_value
|
|
mock_collection.insert_one.reset_mock()
|
|
yield mock_collection.insert_one
|
|
|
|
|
|
@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='laborious',
|
|
)
|
|
yield handler
|
|
handler.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def metrics_controller(mock_logger):
|
|
"""Create a real MetricsController instance."""
|
|
return MetricsController(logger=mock_logger)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_minio_repository():
|
|
"""Mock MinIO repository for object storage operations."""
|
|
mock_repo = MagicMock()
|
|
|
|
# Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve()
|
|
# can decode the payload if offloading is exercised in an integration scenario.
|
|
parquet_df = pd.DataFrame({'a': [1]})
|
|
parquet_buffer = BytesIO()
|
|
parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
|
|
parquet_bytes = parquet_buffer.getvalue()
|
|
|
|
# sientia_do MinioRepository API
|
|
mock_repo.bucket = 'test-bucket'
|
|
mock_repo.upload_file = AsyncMock(
|
|
side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: {
|
|
'minio_object_name': f'sientia/streamlit-connectors/{relative_key}',
|
|
'original_filename': relative_key.rsplit('/', 1)[-1],
|
|
'uploaded_at': '2024-01-01T00:00:00Z',
|
|
'sha256_hash': 'deadbeef',
|
|
}
|
|
)
|
|
mock_repo.download_file = AsyncMock(return_value=parquet_bytes)
|
|
mock_repo.list_objects = AsyncMock(return_value=[])
|
|
mock_repo.delete_file = AsyncMock()
|
|
mock_repo.close = MagicMock()
|
|
|
|
return mock_repo
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_pi_web_api_repository():
|
|
"""Mock PI Web API repository for PI Web API operations."""
|
|
mock_repo = MagicMock()
|
|
|
|
async def _write_value(web_ids, value, metadata=None, **kwargs):
|
|
"""
|
|
Mirror successful PI writes: one response item per requested web_id.
|
|
|
|
write_pi_web_api_data passes the list into process_pi_web_api_response (not a
|
|
wrapped {'Items': ...} envelope).
|
|
"""
|
|
return [{'WebId': wid, 'Errors': []} for wid in web_ids]
|
|
|
|
mock_repo.write_value = AsyncMock(side_effect=_write_value)
|
|
mock_repo.close = MagicMock()
|
|
return mock_repo
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_opc_repository():
|
|
"""Mock OPC repository for OPC operations."""
|
|
mock_repo = MagicMock()
|
|
mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
|
|
mock_repo.disconnect = AsyncMock()
|
|
return mock_repo
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def opc_e2e_server():
|
|
"""
|
|
In-process asyncua OPC UA server for E2E tests against OpcRepository.
|
|
"""
|
|
server = OpcE2ETestServer()
|
|
await server.start()
|
|
try:
|
|
yield server
|
|
finally:
|
|
await server.stop()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def patch_create_engine(postgres_engine):
|
|
"""Patch create_engine to return test postgres_engine."""
|
|
with patch(
|
|
'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def patch_minio_repository(mock_minio_repository):
|
|
"""Patch MinioRepository to return mock."""
|
|
# Patch where Activities resolves the symbol (import binds the original class).
|
|
with patch(
|
|
'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def patch_pi_web_api_repository(mock_pi_web_api_repository):
|
|
"""Patch MLflowRepository to return mock."""
|
|
with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def mock_mlflow_models():
|
|
"""Create mock models for MLflow load_model methods."""
|
|
# Mock transform model - returns DataFrame with same index as input
|
|
mock_transform_model = MagicMock()
|
|
|
|
def mock_transform_predict(data):
|
|
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
|
print(data.to_csv())
|
|
print(data.index)
|
|
result = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.234] * num_rows,
|
|
'feature_2': [0.783] * num_rows,
|
|
}
|
|
)
|
|
result.index = data.index
|
|
return result
|
|
|
|
mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
|
|
|
|
# Mock predict model - returns array/list of predictions
|
|
mock_predict_model = MagicMock()
|
|
|
|
def mock_predict_predict(data):
|
|
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
|
return [0.5] * num_rows
|
|
|
|
mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
|
|
|
|
# Mock PyFuncModel for compressed models
|
|
mock_pyfunc_model = MagicMock()
|
|
mock_pyfunc_model._model_impl = MagicMock()
|
|
mock_pyfunc_model._model_impl.python_model = mock_transform_model
|
|
|
|
return {
|
|
'transform_model': mock_transform_model,
|
|
'predict_model': mock_predict_model,
|
|
'pyfunc_model': mock_pyfunc_model,
|
|
}
|
|
|
|
|
|
def _build_joblib_fallback_estimator():
|
|
"""
|
|
Build a real sklearn estimator for the joblib fallback artifact.
|
|
|
|
`request_predict` appends a string `timestamp` column to the transformed
|
|
data, so the estimator selects the transform output columns explicitly and
|
|
drops the rest.
|
|
"""
|
|
features = ['feature_1', 'feature_2']
|
|
estimator = Pipeline(
|
|
steps=[
|
|
(
|
|
'select',
|
|
ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
|
|
),
|
|
('regressor', LinearRegression()),
|
|
]
|
|
)
|
|
train = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.0, 1.0, 2.0, 3.0],
|
|
'feature_2': [0.0, 1.0, 0.0, 1.0],
|
|
'timestamp': ['2024-01-01 12:00:00'] * 4,
|
|
}
|
|
)
|
|
estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
|
|
return estimator
|
|
|
|
|
|
@pytest.fixture
|
|
def joblib_fallback_artifact(tmp_path):
|
|
"""
|
|
Real artifact directory holding a joblib-compressed sklearn estimator.
|
|
|
|
Mirrors what the joblib fallback finds after downloading `prediction_model`:
|
|
a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
|
|
`MLmodel` placeholder.
|
|
|
|
Return:
|
|
dict with 'path' (the artifact directory) and 'estimator' (the object
|
|
that was serialized, so tests can compute the expected prediction).
|
|
"""
|
|
artifact_dir = tmp_path / 'prediction_model'
|
|
artifact_dir.mkdir()
|
|
|
|
estimator = _build_joblib_fallback_estimator()
|
|
joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
|
|
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
|
|
|
|
return {'path': str(artifact_dir), 'estimator': estimator}
|
|
|
|
|
|
@pytest.fixture
|
|
def joblib_fallback_artifact_without_pickle(tmp_path):
|
|
"""Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
|
|
artifact_dir = tmp_path / 'prediction_model_without_pickle'
|
|
artifact_dir.mkdir()
|
|
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
|
|
return str(artifact_dir)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
|
|
"""Patch mlflow module in repository with load_model mocks."""
|
|
mock_mlflow = MagicMock()
|
|
|
|
# Mock sklearn.load_model
|
|
def mock_sklearn_load_model(model_uri):
|
|
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
|
|
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
|
|
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
|
return mock_mlflow_models['transform_model']
|
|
return mock_mlflow_models['predict_model']
|
|
|
|
mock_mlflow.sklearn = MagicMock()
|
|
mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
|
|
|
|
# Mock pyfunc.load_model
|
|
def mock_pyfunc_load_model(model_uri):
|
|
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
|
|
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
|
|
if 'artifacts' in model_uri or 'tmp' in model_uri:
|
|
return mock_mlflow_models['pyfunc_model']
|
|
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
|
return mock_mlflow_models['transform_model']
|
|
return mock_mlflow_models['predict_model']
|
|
|
|
mock_mlflow.pyfunc = MagicMock()
|
|
mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
|
|
|
|
# Mock pytorch.load_model
|
|
mock_mlflow.pytorch = MagicMock()
|
|
mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
|
|
|
|
# Mock other mlflow methods that might be called
|
|
mock_mlflow.set_tracking_uri = MagicMock()
|
|
mock_mlflow.get_run = MagicMock(
|
|
return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
|
|
)
|
|
mock_mlflow.tracking = MagicMock()
|
|
mock_mlflow_client = MagicMock(
|
|
search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
|
|
search_model_versions=MagicMock(
|
|
return_value=[
|
|
MagicMock(
|
|
current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
|
|
)
|
|
]
|
|
),
|
|
)
|
|
# Only the joblib fallback calls download_artifacts directly on the client.
|
|
mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
|
|
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
|
|
|
|
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
|
|
yield mock_mlflow
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_client(patch_mlflow):
|
|
"""The MlflowClient mock the repository holds as `self.client`."""
|
|
return patch_mlflow.tracking.MlflowClient.return_value
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def test_activities(
|
|
postgres_engine,
|
|
postgres_container,
|
|
mock_logger,
|
|
notification_handler,
|
|
metrics_controller,
|
|
mock_minio_repository,
|
|
patch_create_engine,
|
|
patch_minio_repository,
|
|
patch_mlflow,
|
|
patch_pi_web_api_repository,
|
|
mock_opc_repository,
|
|
):
|
|
"""
|
|
Create Activities instance with test dependencies.
|
|
|
|
This fixture creates a real Activities instance with:
|
|
- PostgreSQL database (via testcontainers)
|
|
- Mocked MinIO client
|
|
- Real NotificationHandler and MetricsController (with mocked underlying services)
|
|
"""
|
|
activities = Activities(
|
|
postgres_config={
|
|
'host': 'localhost',
|
|
'port': postgres_container.get_exposed_port(5432),
|
|
'user': 'test',
|
|
'password': 'test',
|
|
'dbname': 'test',
|
|
'min_connections': 1,
|
|
'max_connections': 5,
|
|
},
|
|
mlflow_config={
|
|
'host': 'http://localhost',
|
|
'port': '5000',
|
|
'username': 'test',
|
|
'password': 'test',
|
|
},
|
|
minio_config={
|
|
# Host:port only; Minio() prepends http(s):// from the secure flag.
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'test',
|
|
'secret_key': 'test',
|
|
'default_bucket': 'test-bucket',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
},
|
|
opc_config={},
|
|
pi_web_api_config={
|
|
'base_url': 'http://localhost:8080',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
},
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
activities.opc_repository = {
|
|
'1': mock_opc_repository,
|
|
}
|
|
|
|
try:
|
|
yield activities
|
|
finally:
|
|
# Cleanup - ALWAYS runs, even if test fails
|
|
await activities.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def test_activities_real_minio(
|
|
postgres_engine,
|
|
postgres_container,
|
|
minio_container,
|
|
mock_logger,
|
|
notification_handler,
|
|
metrics_controller,
|
|
patch_create_engine,
|
|
patch_mlflow,
|
|
patch_pi_web_api_repository,
|
|
mock_opc_repository,
|
|
):
|
|
"""
|
|
Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests.
|
|
"""
|
|
minio_client = minio_container.get_client()
|
|
if not minio_client.bucket_exists('test-bucket'):
|
|
minio_client.make_bucket('test-bucket')
|
|
minio_port = minio_container.get_exposed_port(9000)
|
|
activities = Activities(
|
|
postgres_config={
|
|
'host': 'localhost',
|
|
'port': postgres_container.get_exposed_port(5432),
|
|
'user': 'test',
|
|
'password': 'test',
|
|
'dbname': 'test',
|
|
'min_connections': 1,
|
|
'max_connections': 5,
|
|
},
|
|
mlflow_config={
|
|
'host': 'http://localhost',
|
|
'port': '5000',
|
|
'username': 'test',
|
|
'password': 'test',
|
|
},
|
|
minio_config={
|
|
'endpoint_url': f'localhost:{minio_port}',
|
|
'access_key': 'minioadmin',
|
|
'secret_key': 'minioadmin',
|
|
'default_bucket': 'test-bucket',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
},
|
|
opc_config={},
|
|
pi_web_api_config={
|
|
'base_url': 'http://localhost:8080',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
},
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
activities.opc_repository = {'1': mock_opc_repository}
|
|
try:
|
|
yield activities
|
|
finally:
|
|
await activities.shutdown()
|
|
|
|
|
|
def _worker_activity_list(test_activities: Activities):
|
|
return [
|
|
test_activities.load_custom_query,
|
|
test_activities.load_query_with_minio_offload,
|
|
test_activities.cleanup_minio_objects_expired,
|
|
test_activities.input_gate,
|
|
test_activities.request_transform,
|
|
test_activities.mlflow_response_gate,
|
|
test_activities.mlflow_content_gate,
|
|
test_activities.request_predict,
|
|
test_activities.repeat_last_prediction,
|
|
test_activities.format_prediction,
|
|
test_activities.format_transformed_data,
|
|
test_activities.format_default_prediction,
|
|
test_activities.write_pi_web_api_data,
|
|
test_activities.write_opc_data,
|
|
test_activities.export_data_to_postgres,
|
|
test_activities.export_payload_to_postgres,
|
|
test_activities.write_metrics,
|
|
]
|
|
|
|
|
|
@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=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
|
activities=_worker_activity_list(test_activities),
|
|
) as worker:
|
|
yield worker
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
|
|
"""Temporal worker backed by Activities using real MinIO testcontainer."""
|
|
async with Worker(
|
|
temporal_test_env.client,
|
|
task_queue='test-queue',
|
|
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
|
activities=_worker_activity_list(test_activities_real_minio),
|
|
) as worker:
|
|
yield worker
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def test_activities_real_opc(
|
|
postgres_engine,
|
|
postgres_container,
|
|
opc_e2e_server: OpcE2ETestServer,
|
|
mock_logger,
|
|
notification_handler,
|
|
metrics_controller,
|
|
patch_create_engine,
|
|
patch_minio_repository,
|
|
patch_mlflow,
|
|
patch_pi_web_api_repository,
|
|
):
|
|
"""
|
|
Activities with a real OpcRepository connected to the in-process OPC UA server.
|
|
"""
|
|
activities = Activities(
|
|
postgres_config={
|
|
'host': 'localhost',
|
|
'port': postgres_container.get_exposed_port(5432),
|
|
'user': 'test',
|
|
'password': 'test',
|
|
'dbname': 'test',
|
|
'min_connections': 1,
|
|
'max_connections': 5,
|
|
},
|
|
mlflow_config={
|
|
'host': 'http://localhost',
|
|
'port': '5000',
|
|
'username': 'test',
|
|
'password': 'test',
|
|
},
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'test',
|
|
'secret_key': 'test',
|
|
'default_bucket': 'test-bucket',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
},
|
|
opc_config={
|
|
'1': {
|
|
'id': '1',
|
|
'server_name': 'e2e-opc',
|
|
'url': opc_e2e_server.url,
|
|
'server_uri': opc_e2e_server.url,
|
|
'cert_path': None,
|
|
'private_key_path': None,
|
|
'server_cert_path': None,
|
|
'reconnection_interval': 0,
|
|
}
|
|
},
|
|
pi_web_api_config={
|
|
'base_url': 'http://localhost:8080',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
},
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
await activities.init_opc()
|
|
repo = activities.opc_repository['1']
|
|
assert repo._session_ready.is_set(), 'OPC E2E server connection failed during init_opc'
|
|
try:
|
|
yield activities
|
|
finally:
|
|
await activities.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture(scope='function')
|
|
async def temporal_worker_real_opc(temporal_test_env, test_activities_real_opc):
|
|
"""Temporal worker backed by Activities using the in-process OPC UA server."""
|
|
async with Worker(
|
|
temporal_test_env.client,
|
|
task_queue='test-queue',
|
|
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
|
activities=_worker_activity_list(test_activities_real_opc),
|
|
) as worker:
|
|
yield worker
|