def test_custom_model_runner_does_not_exist___default_runner_is_used(self):
with TemporaryDirectory() as inputs_dir, \
TemporaryDirectory() as model_data_dir, \
TemporaryDirectory() as work_dir, \
TemporaryDirectory() as out_dir:
with SettingsPatcher(INPUTS_DATA_DIRECTORY=inputs_dir,
MODEL_DATA_DIRECTORY=model_data_dir,
WORKING_DIRECTORY=work_dir,
OUTPUTS_DATA_DIRECTORY=out_dir):
self.create_tar(str(Path(inputs_dir, 'location.tar')))
Path(model_data_dir, 'supplier', 'version').mkdir(parents=True)
with patch('src.model_execution_worker.tasks.runner'
) as default_mock:
start_analysis(
{
'analysis_settings': {
'source_tag': 'source',
'analysis_tag': 'source',
'module_supplier_id': 'supplier',
'model_version_id': 'version'
}
},
'location',
)
default_mock.run.assert_called_once_with(
{
'source_tag': 'source',
'analysis_tag': 'source',
'module_supplier_id': 'supplier',
'model_version_id': 'version'
}, settings.getint('worker', 'KTOOLS_BATCH_COUNT'))
def test_do_clear_working_is_false___working_directory_is_not_removed_after_run(
self):
with TemporaryDirectory() as inputs_dir, \
TemporaryDirectory() as model_data_dir, \
TemporaryDirectory() as work_dir, \
TemporaryDirectory() as out_dir:
with SettingsPatcher(INPUTS_DATA_DIRECTORY=inputs_dir,
MODEL_DATA_DIRECTORY=model_data_dir,
WORKING_DIRECTORY=work_dir,
OUTPUTS_DATA_DIRECTORY=out_dir,
DO_CLEAr_WORKING='False'):
self.create_tar(str(Path(inputs_dir, 'location.tar')))
Path(model_data_dir, 'supplier', 'version').mkdir(parents=True)
with patch('src.model_execution_worker.tasks.runner'):
start_analysis(
{
'analysis_settings': {
'source_tag': 'source',
'analysis_tag': 'source',
'module_supplier_id': 'supplier',
'model_version_id': 'version'
}
},
'location',
)
self.assertGreater(len(os.listdir(work_dir)), 0)
def test_settings_file_does_not_exist___exception_is_raised(self):
with TemporaryDirectory() as media_root:
with SettingsPatcher(MEDIA_ROOT=media_root):
self.create_tar(str(Path(media_root, 'location.tar')))
with self.assertRaises(MissingInputsException):
start_analysis(input_location=os.path.join(
media_root, 'location.tar'),
analysis_settings=os.path.join(
media_root, 'analysis_settings.json'))
def test_custom_model_runner_does_not_exist___generate_losses_is_called_output_files_are_tared_up(
self):
with TemporaryDirectory() as media_root, \
TemporaryDirectory() as model_data_dir, \
TemporaryDirectory() as run_dir, \
TemporaryDirectory() as work_dir:
with SettingsPatcher(
MODEL_SUPPLIER_ID='supplier',
MODEL_ID='model',
MODEL_VERSION_ID='version',
MEDIA_ROOT=media_root,
MODEL_DATA_DIRECTORY=model_data_dir,
WORKING_DIRECTORY=work_dir,
):
self.create_tar(str(Path(media_root, 'location.tar')))
Path(media_root, 'analysis_settings.json').touch()
Path(run_dir, 'output').mkdir(parents=True)
Path(model_data_dir, 'supplier', 'model',
'version').mkdir(parents=True)
cmd_instance = Mock()
cmd_instance.stdout = b'output'
cmd_instance.stderr = b'errors'
@contextmanager
def fake_run_dir(*args, **kwargs):
yield run_dir
with patch('src.model_execution_worker.tasks.subprocess.run', Mock(return_value=cmd_instance)) as cmd_mock, \
patch('src.model_execution_worker.tasks.get_worker_versions', Mock(return_value='')), \
patch('src.model_execution_worker.tasks.filestore.compress') as tarfile, \
patch('src.model_execution_worker.tasks.TemporaryDir', fake_run_dir):
output_location, log_location, error_location, returncode = start_analysis(
os.path.join(media_root, 'analysis_settings.json'),
os.path.join(media_root, 'location.tar'),
)
test_env = os.environ.copy()
cmd_mock.assert_called_once_with([
'oasislmf',
'model',
'generate-losses',
'--oasis-files-dir',
os.path.join(run_dir, 'input'),
'--config',
get_oasislmf_config_path(
settings.get('worker', 'model_id')),
'--model-run-dir',
run_dir,
'--analysis-settings-json',
os.path.join(media_root, 'analysis_settings.json'),
'--ktools-fifo-relative',
'--verbose',
],
stderr=subprocess.PIPE,
stdout=subprocess.PIPE,
env=test_env)
tarfile.assert_called_once_with(
output_location, os.path.join(run_dir, 'output'),
'output')
def test_custom_model_runner_exists___custom_runner_is_used(self):
with TemporaryDirectory() as inputs_dir, \
TemporaryDirectory() as model_data_dir, \
TemporaryDirectory() as work_dir, \
TemporaryDirectory() as out_dir, \
TemporaryDirectory() as module_dir:
with SettingsPatcher(INPUTS_DATA_DIRECTORY=inputs_dir,
MODEL_DATA_DIRECTORY=model_data_dir,
WORKING_DIRECTORY=work_dir,
OUTPUTS_DATA_DIRECTORY=out_dir,
SUPPLIER_MODULE_DIRECTORY=module_dir):
self.create_tar(str(Path(inputs_dir, 'location.tar')))
Path(model_data_dir, 'supplier', 'version').mkdir(parents=True)
Path(module_dir, 'supplier').mkdir()
Path(module_dir, 'supplier', '__init__.py').touch()
with open(
str(
Path(module_dir, 'supplier').joinpath(
'supplier_model_runner.py')), 'w') as module:
module.writelines([
'from pathlib2 import Path\n',
'def run(settings, location):\n',
' Path("{}", "custom_model").touch()\n'.format(
out_dir)
])
start_analysis(
{
'analysis_settings': {
'source_tag': 'source',
'analysis_tag': 'source',
'module_supplier_id': 'supplier',
'model_version_id': 'version'
}
},
'location',
)
self.assertTrue(Path(out_dir, "custom_model").exists())
def test_custom_model_runner_does_not_exist___generate_losses_is_called_output_files_are_tared_up(
self):
with TemporaryDirectory() as media_root, \
TemporaryDirectory() as model_data_dir, \
TemporaryDirectory() as work_dir:
with SettingsPatcher(
MODEL_SUPPLIER_ID='supplier',
MODEL_ID='model',
MODEL_VERSION_ID='version',
MEDIA_ROOT=media_root,
MODEL_DATA_DIRECTORY=model_data_dir,
WORKING_DIRECTORY=work_dir,
):
self.create_tar(str(Path(media_root, 'location.tar')))
Path(model_data_dir, 'supplier', 'model',
'version').mkdir(parents=True)
cmd_instance = Mock()
with patch('src.model_execution_worker.tasks.GenerateLossesCmd', Mock(return_value=cmd_instance)) as cmd_mock, \
patch('src.model_execution_worker.tasks.tarfile') as tarfile:
output_location = start_analysis(
'analysis_settings.json',
'location.tar',
)
cmd_mock.assert_called_once_with(argv=[
'--oasis-files-dir', ANY, '--config',
get_oasislmf_config_path(
settings.get('worker', 'model_id')),
'--model-run-dir', ANY, '--analysis-settings-json',
'analysis_settings.json', '--ktools-num-processes',
settings.get('worker', 'KTOOLS_BATCH_COUNT'),
'--ktools-alloc-rule-gul',
settings.get('worker', 'KTOOLS_ALLOC_RULE_GUL'),
'--ktools-alloc-rule-il',
settings.get('worker', 'KTOOLS_ALLOC_RULE_IL'),
'--ktools-alloc-rule-ri',
settings.get('worker', 'KTOOLS_ALLOC_RULE_RI'),
'--ktools-fifo-relative'
])
cmd_instance.run.assert_called_once_with()
self.assertEqual(
tarfile.open.call_args_list[1][0],
(str(Path(media_root, output_location)), 'w:gz'))
