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.
212 lines
6.7 KiB
Python
212 lines
6.7 KiB
Python
"""
|
|
E2E tests for OPC export using an in-process asyncua server and real OpcRepository.
|
|
|
|
Covers scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 from e2e/scenarios.md.
|
|
Mock-based OPC tests remain in test_predictions_batch_format_export.py.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import (
|
|
assert_prediction,
|
|
insert_sample_data,
|
|
make_workflow_id,
|
|
start_and_await_workflow,
|
|
)
|
|
from e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
|
|
from e2e.test_predictions_batch_format_export import get_base_input_data
|
|
from laborious.activities.activities import Activities
|
|
from laborious.activities.opc import OPC_RECONNECT_IN_PROGRESS_COMMENT
|
|
from laborious.utils.repository.opc_repository import OpcRepository
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
|
|
async def _reconnect_under_held_lock(
|
|
repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
|
|
) -> None:
|
|
"""
|
|
Hold the connection lock until released, then reconnect.
|
|
|
|
The lock must stay held for the whole workflow, and how long the workflow takes to
|
|
reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked`
|
|
finish first, which sets `_session_ready` again and the write then succeeds.
|
|
|
|
Args:
|
|
repo (OpcRepository): Connected repository.
|
|
acquired (asyncio.Event): Set once the lock is held, before any write can race.
|
|
release (asyncio.Event): Awaited while holding the lock; set by the caller.
|
|
"""
|
|
async with repo._connection_lock:
|
|
acquired.set()
|
|
await release.wait()
|
|
await repo._reconnect_locked()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
@pytest.mark.opc
|
|
async def test_scenario_3_1_2_export_with_opc_only_real_server(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_real_opc: Worker,
|
|
test_activities_real_opc: Activities,
|
|
opc_e2e_server: OpcE2ETestServer,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 3.1.2 (real OPC): connect, write prediction and confidence, verify server values.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 412
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
|
input_data['pi_web_api_output_config'] = None
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-opc-real-happy'),
|
|
)
|
|
|
|
test_activities_real_opc.pi_web_api_client.write_value.assert_not_called()
|
|
assert await opc_e2e_server.read_prediction() == pytest.approx(0.5)
|
|
assert await opc_e2e_server.read_confidence() == pytest.approx(0.0)
|
|
assert_prediction(postgres_engine, model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
@pytest.mark.opc
|
|
async def test_scenario_3_2_2_opc_write_error_real_server(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_real_opc: Worker,
|
|
test_activities_real_opc: Activities,
|
|
opc_e2e_server: OpcE2ETestServer,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 3.2.2 (real OPC): unknown NodeId yields generic write failure (confidence 12).
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 422
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
|
|
node_ids = opc_e2e_server.node_ids
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['opc_output_config'] = build_opc_output_config(
|
|
node_ids,
|
|
prediction_tag=UNKNOWN_NODE_ID,
|
|
confidence_tag=UNKNOWN_NODE_ID,
|
|
)
|
|
input_data['pi_web_api_output_config'] = None
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-opc-real-bad-node'),
|
|
)
|
|
|
|
assert_prediction(
|
|
postgres_engine,
|
|
model_id,
|
|
prediction_confidence=12,
|
|
comments='Some data could not be written to OPC servers',
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
@pytest.mark.opc
|
|
async def test_scenario_3_2_4_opc_session_bad_real_server(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_real_opc: Worker,
|
|
test_activities_real_opc: Activities,
|
|
opc_e2e_server: OpcE2ETestServer,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 3.2.4 (real OPC): server PreWrite fault injects BadSessionIdInvalid (confidence 14).
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 424
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
|
|
opc_e2e_server.set_session_bad_on_write(True)
|
|
try:
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['opc_output_config'] = build_opc_output_config(
|
|
opc_e2e_server.node_ids,
|
|
prediction_only=True,
|
|
)
|
|
input_data['pi_web_api_output_config'] = None
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-opc-real-session-bad'),
|
|
)
|
|
finally:
|
|
opc_e2e_server.set_session_bad_on_write(False)
|
|
|
|
assert_prediction(
|
|
postgres_engine,
|
|
model_id,
|
|
prediction_confidence=14,
|
|
comments_contains='OPC UA session/channel error: BadSessionIdInvalid',
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
@pytest.mark.opc
|
|
async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_real_opc: Worker,
|
|
test_activities_real_opc: Activities,
|
|
opc_e2e_server: OpcE2ETestServer,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 3.2.5 (real OPC): writes rejected while reconnect holds the connection lock.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 425
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
|
|
repo = test_activities_real_opc.opc_repository['1']
|
|
repo._session_ready.clear()
|
|
acquired = asyncio.Event()
|
|
release = asyncio.Event()
|
|
reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
|
|
await acquired.wait()
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
|
input_data['pi_web_api_output_config'] = None
|
|
|
|
try:
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-opc-real-reconnect-block'),
|
|
)
|
|
finally:
|
|
release.set()
|
|
await reconnect_task
|
|
|
|
assert_prediction(
|
|
postgres_engine,
|
|
model_id,
|
|
prediction_confidence=14,
|
|
comments_contains=OPC_RECONNECT_IN_PROGRESS_COMMENT,
|
|
)
|