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'))