vitor-aignosi
dfb9deae1a
QTZPOC-20: Remove PyTorch from supported model flavors as it was never functional due to missing runtime dependencies.
845 lines
29 KiB
Python
845 lines
29 KiB
Python
"""
|
|
Pytest configuration and fixtures for E2E tests.
|
|
"""
|
|
|
|
import importlib
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
# Load every BSON codec now, for its side effect only. `pymongo` imports `bson.objectid`,
|
|
# `bson.timestamp`, `bson.decimal128` and a dozen more lazily, and the first time is while an
|
|
# activity is running — after Temporal's workflow sandbox has snapshotted `sys.modules`. The sandbox
|
|
# then warns once per submodule ("Module bson.timestamp was imported after initial workflow load"),
|
|
# which was ~20 lines of noise per test that touches MongoDB, enough to bury a warning worth reading.
|
|
# This has to happen **before** the workflow modules are imported below, which is why it is here and
|
|
# not in `import_harness.py` — pytest loads `pytest_plugins` only after this module's body has run.
|
|
import bson.json_util # noqa: F401
|
|
|
|
# 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 sqlalchemy import create_engine, text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
from testcontainers.community.minio import MinioContainer
|
|
from testcontainers.community.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
|
|
|
|
# The `.sientia` import fixtures live in their own module rather than here: four containers' worth
|
|
# of setup, an MLflow server subprocess and an ImportModel worker would double this file. Every
|
|
# fixture it adds is `import_`-prefixed, so nothing below changes meaning.
|
|
pytest_plugins = ['e2e.import_harness']
|
|
|
|
# Test constants
|
|
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
|
|
TEST_DATABASE_NAME = 'test_db'
|
|
|
|
# Dedicated model served through the explicit `joblib` flavor, so no MLflow flavor loader runs.
|
|
JOBLIB_FLAVOR_MODEL_NAME = 'joblib_flavor_model'
|
|
|
|
# Module carrying the classes of the joblib artifacts' models. It ships inside the artifact's
|
|
# `code/` directory and is importable only while that directory is on `sys.path` — which is what
|
|
# the joblib flavor arranges before deserializing.
|
|
JOBLIB_FLAVOR_CODE_MODULE = 'joblib_flavor_probe'
|
|
|
|
JOBLIB_FLAVOR_WEIGHTS = {'feature_1': 2.0, 'feature_2': 3.0}
|
|
|
|
JOBLIB_FLAVOR_CODE_SOURCE = '''"""Model classes that exist only inside the artifact's `code/` directory."""
|
|
|
|
import pandas as pd
|
|
|
|
FEATURES = ('feature_1', 'feature_2')
|
|
|
|
|
|
class JoblibFlavorPredictor:
|
|
"""Prediction model: a weighted sum of the transform output columns."""
|
|
|
|
def __init__(self, weights):
|
|
self.weights = weights
|
|
|
|
def fit(self, data):
|
|
return self
|
|
|
|
def predict(self, data):
|
|
frame = pd.DataFrame(data)
|
|
return [
|
|
sum(row[name] * self.weights[name] for name in FEATURES)
|
|
for _, row in frame.iterrows()
|
|
]
|
|
|
|
|
|
class JoblibFlavorTransformer:
|
|
"""Transform model: emits the fixed feature frame the prediction model expects."""
|
|
|
|
target_variable = 'target'
|
|
|
|
def fit(self, data):
|
|
return self.predict(data)
|
|
|
|
def predict(self, data):
|
|
frame = pd.DataFrame(data)
|
|
result = pd.DataFrame(
|
|
{'feature_1': [0.234] * len(frame), 'feature_2': [0.783] * len(frame)}
|
|
)
|
|
result.index = frame.index
|
|
return result
|
|
'''
|
|
|
|
|
|
@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.
|
|
|
|
`log_statement=all` is the one thing the import suite needs from this fixture: proving an import
|
|
issues no statement at all against the predictions database requires a log that would have
|
|
shown one. It changes what the server writes to its own log and nothing a test observes.
|
|
"""
|
|
postgres = PostgresContainer('postgres:15').with_command('postgres -c log_statement=all')
|
|
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 _write_artifact_code_dir(artifact_dir):
|
|
"""Write the artifact's `code/` directory, carrying the module with the model classes."""
|
|
code_dir = artifact_dir / 'code'
|
|
code_dir.mkdir()
|
|
(code_dir / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').write_text(JOBLIB_FLAVOR_CODE_SOURCE)
|
|
return code_dir
|
|
|
|
|
|
def _build_joblib_flavor_artifact(artifact_dir, class_name, *args):
|
|
"""
|
|
Build one joblib artifact: `code/` with the class module, plus the compressed `model.pkl`.
|
|
|
|
Instantiating **and** dumping happen while `code/` is on `sys.path` — `pickle` refuses to
|
|
serialize a class it cannot import. Both are undone on the way out, so from here the pickled
|
|
class is unreachable until something puts the artifact's `code/` back on the path. That
|
|
something is the flavor under test.
|
|
"""
|
|
code_dir = _write_artifact_code_dir(artifact_dir)
|
|
|
|
sys.path.insert(0, str(code_dir))
|
|
try:
|
|
module = importlib.import_module(JOBLIB_FLAVOR_CODE_MODULE)
|
|
model = getattr(module, class_name)(*args)
|
|
joblib.dump(model, artifact_dir / 'model.pkl', compress=3)
|
|
finally:
|
|
sys.path.remove(str(code_dir))
|
|
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
|
importlib.invalidate_caches()
|
|
|
|
return {'path': str(artifact_dir), 'code_dir': str(code_dir), 'model': model}
|
|
|
|
|
|
@pytest.fixture
|
|
def joblib_flavor_artifacts(tmp_path):
|
|
"""
|
|
Real artifact directories for the `joblib` flavor, one per model type.
|
|
|
|
Each holds a top-level `model.pkl` written by `joblib.dump(..., compress=3)` and a real
|
|
`code/` directory with the module where the model's class lives. No `MLmodel` is written:
|
|
the artifact this flavor reads — and the one a joblib retrain writes — has none.
|
|
|
|
Return:
|
|
dict keyed by artifact name ('prediction_model', 'data_model'), each value carrying
|
|
'path', 'code_dir' and 'model' (the object serialized, so tests can compute the
|
|
expected prediction).
|
|
"""
|
|
original_sys_path = list(sys.path)
|
|
|
|
builds = {
|
|
'prediction_model': ('JoblibFlavorPredictor', (JOBLIB_FLAVOR_WEIGHTS,)),
|
|
'data_model': ('JoblibFlavorTransformer', ()),
|
|
}
|
|
|
|
artifacts = {}
|
|
|
|
for artifact_name, (class_name, args) in builds.items():
|
|
artifact_dir = tmp_path / artifact_name
|
|
artifact_dir.mkdir()
|
|
|
|
artifacts[artifact_name] = _build_joblib_flavor_artifact(artifact_dir, class_name, *args)
|
|
|
|
yield artifacts
|
|
|
|
# The flavor prepends `code/` to `sys.path` and deliberately never removes it.
|
|
sys.path[:] = original_sys_path
|
|
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
|
importlib.invalidate_caches()
|
|
|
|
|
|
@pytest.fixture
|
|
def joblib_flavor_artifact_without_pickle(tmp_path):
|
|
"""Artifact directory with no top-level `.pkl`, so the joblib flavor fails on it."""
|
|
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_flavor_artifacts):
|
|
"""Patch mlflow module in repository with load_model mocks."""
|
|
mock_mlflow = MagicMock()
|
|
|
|
# Mock sklearn.load_model
|
|
def mock_sklearn_load_model(model_uri):
|
|
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 '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 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 flavor calls download_artifacts directly on the client, one artifact
|
|
# per model type — a single return value would serve `data_model` from `prediction_model`.
|
|
def mock_download_artifacts(_run_id, artifact_name):
|
|
return joblib_flavor_artifacts[artifact_name]['path']
|
|
|
|
mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts)
|
|
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
|