vitor-aignosi
30ff378cb7
QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.
861 lines
30 KiB
Python
861 lines
30 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_mlflow.artifacts = MagicMock()
|
|
|
|
# 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'
|
|
)
|
|
]
|
|
),
|
|
)
|
|
|
|
# `download_model` asks whether the run's artifact carries a `code` directory before it loads
|
|
# anything, and downloads only that directory when it does. Answering "no" here keeps the
|
|
# download counts of these tests about the model itself; a test that exercises the `code` step
|
|
# overrides `list_artifacts` to answer "yes".
|
|
def mock_list_artifacts(_run_id, artifact_path=None):
|
|
if artifact_path:
|
|
return []
|
|
return [MagicMock(path=artifact_name) for artifact_name in joblib_flavor_artifacts]
|
|
|
|
mock_mlflow_client.list_artifacts = MagicMock(side_effect=mock_list_artifacts)
|
|
|
|
# The joblib flavor downloads a whole model artifact, one per model type — a single return
|
|
# value would serve `data_model` from `prediction_model`. The `code` step asks for the nested
|
|
# `<artifact>/code` instead, and gets that directory back, as MLflow would return it.
|
|
def mock_download_artifacts(_run_id, artifact_name):
|
|
if artifact_name.endswith('/code'):
|
|
return joblib_flavor_artifacts[artifact_name.removesuffix('/code')]['code_dir']
|
|
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
|