async def test_simple_record_transmit_from_file(
record_transmitter_factory_context: ContextManager[Callable[
[str], RecordTransmitter]],
record_in,
expected_data,
mime_type,
storage_path,
):
if isinstance(record_in, (NumericalRecordTree, BlobRecordTree)):
pytest.skip("unsupported serialization of opaque RecordTree")
if isinstance(record_in,
BlobRecord) and mime_type != "application/octet-stream":
pytest.skip(
f"unsupported serialization of opaque record to {mime_type}")
if (isinstance(record_in, NumericalRecord)
and mime_type == "application/octet-stream"):
pytest.skip(f"unsupported serialization of num record to {mime_type}")
filename = "record.file"
with record_transmitter_factory_context(
storage_path=storage_path) as record_transmitter_factory, tmp():
transmitter = record_transmitter_factory(name="some_name")
if mime_type == "application/octet-stream":
async with aiofiles.open(filename, "wb") as fb:
await fb.write(expected_data)
else:
await ert.serialization.get_serializer(mime_type).encode_to_path(
expected_data, filename)
transformation = ert.data.FileRecordTransformation()
record = await transformation.transform_output(location=filename,
mime=mime_type)
await transmitter.transmit_record(record)
assert transmitter.is_transmitted()
with pytest.raises(RuntimeError, match="Record already transmitted"):
await transmitter.transmit_record(record)
async def test_simple_record_transmit_and_dump(
record_transmitter_factory_context: ContextManager[Callable[
[str], RecordTransmitter]],
record_in,
expected_data,
mime_type,
storage_path,
):
if isinstance(record_in, (NumericalRecordTree, BlobRecordTree)):
pytest.skip("unsupported serialization of binary RecordTree")
if isinstance(record_in,
BlobRecord) and mime_type != "application/octet-stream":
pytest.skip(
f"unsupported serialization of binary record to {mime_type}")
if (isinstance(record_in, NumericalRecord)
and mime_type == "application/octet-stream"):
pytest.skip(f"unsupported serialization of num record to {mime_type}")
with record_transmitter_factory_context(
storage_path=storage_path) as record_transmitter_factory, tmp():
transmitter = record_transmitter_factory(name="some_name")
await transmitter.transmit_record(record_in)
transformation = SerializationTransformation(
location=pathlib.Path("record"),
mime=mime_type,
direction=ert.data.TransformationDirection.FROM_RECORD,
)
record = await transmitter.load()
await transformation.from_record(record)
if mime_type == "application/octet-stream":
with open("record", "rb") as fb:
assert expected_data == fb.read()
else:
with open("record", "rt", encoding="utf-8") as ft:
assert (ert.serialization.get_serializer(mime_type).encode(
expected_data) == ft.read())
async def test_atomic_to_record_transformation(
record_transmitter_factory_context: ContextManager[
Callable[[str], RecordTransmitter]
],
cls: Type[RecordTransformation],
args: list,
files: List[str],
record_cls: Type[Record],
storage_path,
tmp_path,
):
with record_transmitter_factory_context(
storage_path=storage_path
) as record_transmitter_factory, file_factory_context(
tmp_path
) as file_factory, tmp():
runpath = pathlib.Path(tmp_path)
file_factory(files=files)
transmitter = record_transmitter_factory(name="trans_custom")
assert transmitter.is_transmitted() is False
# TODO: https://github.com/python/mypy/issues/6799
transformation = cls(
*args, direction=TransformationDirection.TO_RECORD
) # type: ignore
record = await transformation.to_record(root_path=runpath)
await transmitter.transmit_record(record)
assert transmitter.is_transmitted()
loaded_record = await transmitter.load()
assert isinstance(loaded_record, record_cls)
async def test_atomic_from_record_transformation(
record_transmitter_factory_context: ContextManager[
Callable[[str], RecordTransmitter]
],
cls: Type[RecordTransformation],
args: list,
type: str,
files: List[str],
storage_path,
tmp_path,
):
runpath = pathlib.Path(".")
with record_transmitter_factory_context(
storage_path=storage_path
) as record_transmitter_factory, record_factory_context(
tmp_path
) as record_factory, tmp():
# TODO: https://github.com/python/mypy/issues/6799
transformation = cls(
*args, direction=TransformationDirection.FROM_RECORD
) # type: ignore
record_in = record_factory(type=type)
transmitter = record_transmitter_factory(name="trans_custom")
await transmitter.transmit_record(record_in)
assert transmitter.is_transmitted()
record = await transmitter.load()
await transformation.from_record(record)
for file in files:
assert (runpath / file).exists()
async def test_atomic_transformation_input(
record_transmitter_factory_context: ContextManager[Callable[
[str], RecordTransmitter]],
transformation_class: RecordTransformation,
transformation_args: list,
location: str,
mime: str,
type: str,
res_files_dumped: List[str],
storage_path,
tmp_path,
):
runpath = pathlib.Path(".")
with record_transmitter_factory_context(
storage_path=storage_path
) as record_transmitter_factory, record_factory_context(
tmp_path) as record_factory, tmp():
transformation = transformation_class(*transformation_args)
if isinstance(transformation, EclSumTransformation):
return # Not supposed to be implemented
record_in = record_factory(type=type)
transmitter = record_transmitter_factory(name="trans_custom")
await transmitter.transmit_record(record_in)
assert transmitter.is_transmitted()
record = await transmitter.load()
await transformation.transform_input(record, mime, runpath,
pathlib.Path(location))
for file in res_files_dumped:
assert (runpath / file).exists()
def prefect_flow_run(ws_monitor, step, input_map, output_map, **kwargs):
"""Use prefect-flow to run task from step output using task_inputs"""
with prefect.context(url=ws_monitor.url, token=None, cert=None):
with Flow("testing") as flow:
task = step.get_task(output_transmitters=output_map,
ee_id="test_ee_id",
**kwargs)
result = task(inputs=input_map)
with tmp():
flow_run = flow.run()
# Stop the mock evaluator WS server
messages = ws_monitor.join_and_get_messages()
return result, flow_run, messages
async def test_overwrite_fail(
record_transmitter_factory_context: ContextManager[
Callable[[str], RecordTransmitter]
],
cls: Type[RecordTransformation],
args: list,
type: str,
files: List[str],
storage_path,
tmp_path,
):
with record_transmitter_factory_context(
storage_path=storage_path
) as record_transmitter_factory, record_factory_context(
tmp_path
) as record_factory, tmp():
if not issubclass(cls, FileTransformation):
pytest.skip(
f"{cls} will not write anything, thus cannot be tested this way"
)
if cls == TarTransformation:
pytest.skip("this protection might not make sense for extracting archives")
transformation = cls(
*args, direction=TransformationDirection.FROM_RECORD
) # type: ignore
record_in = record_factory(type=type)
transmitter = record_transmitter_factory(name="trans_custom")
await transmitter.transmit_record(record_in)
assert transmitter.is_transmitted()
record = await transmitter.load()
if cls == TreeSerializationTransformation:
for file_ in files:
(tmp_path / file_).touch()
elif cls == ExecutableTransformation:
(tmp_path / _BIN_FOLDER).mkdir()
(tmp_path / _BIN_FOLDER / transformation.location).touch()
else:
(tmp_path / transformation.location).touch()
with pytest.raises(FileExistsException):
await transformation.from_record(record, root_path=tmp_path)
def test_on_task_failure(
mock_ws_monitor,
step_test_script_transmitter,
retries,
nfails,
expect,
):
"""Test both job and task failure of prefect-flow-run"""
with tmp() as runpath:
step, input_map, output_map = get_step(
step_name="test_step",
inputs=[(
"script",
Path("unix_test_retry_script.py"),
"application/x-python",
lambda _t=step_test_script_transmitter: _t,
)],
outputs=[],
jobs=[("script", Path("unix_test_retry_script.py"), [runpath])],
type_="unix",
)
result, flow_run, messages = prefect_flow_run(
ws_monitor=mock_ws_monitor,
step=step,
input_map=input_map,
output_map=output_map,
max_retries=retries,
retry_delay=timedelta(seconds=1),
on_failure=functools.partial(_on_task_failure, ee_id="test_ee_id"),
)
task_result = flow_run.result[result]
assert task_result.is_successful() == expect
assert flow_run.is_successful() == expect
fail_job_messages = [
msg for msg in messages if ids.EVTYPE_FM_JOB_FAILURE in msg
]
fail_step_messages = [
msg for msg in messages if ids.EVTYPE_FM_STEP_FAILURE in msg
]
expected_job_failed_messages = 2
assert expected_job_failed_messages == len(fail_job_messages)
assert nfails == len(fail_step_messages)
async def test_atomic_transformation_output(
record_transmitter_factory_context: ContextManager[Callable[
[str], RecordTransmitter]],
transformation_class: RecordTransformation,
transformation_args: list,
location: str,
mime: str,
type: str,
res_files_dumped: List[str],
storage_path,
tmp_path,
):
with record_transmitter_factory_context(
storage_path=storage_path
) as record_transmitter_factory, file_factory_context(
tmp_path) as file_factory, tmp():
runpath = tmp_path
file_factory(files=res_files_dumped)
transmitter = record_transmitter_factory(name="trans_custom")
assert transmitter.is_transmitted() is False
transformation = transformation_class(*transformation_args)
if type != "dir":
location = res_files_dumped[0]
if isinstance(transformation,
RecordTreeTransformation) and not isinstance(
transformation, EclSumTransformation):
return # Not implemented
record = await transformation.transform_output(mime,
runpath / location)
await transmitter.transmit_record(record)
assert transmitter.is_transmitted()
loaded_record = await transmitter.load()
if isinstance(transformation, EclSumTransformation):
assert isinstance(loaded_record, NumericalRecordTree)
else:
assert isinstance(loaded_record, BlobRecord)