def test_migrate_params_retval_does_not_match_schema(self): # LoadedModule.migrate_params() may return invalid data: it's up to the # caller to validate it. In this test, we test that indeed, invalid # data may be returned. code = b"def migrate_params(params):\n return {}" minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/imported.py", Body=code, ContentLength=len(code), ) with self.assertLogs("cjwstate.modules.loaded_module"): lm = LoadedModule.for_module_version( MockModuleVersion( "imported", "abcdef", ParamDType.Dict({"x": ParamDType.String()}), "now", ) ) self.assertEqual( # should have 'x' key lm.migrate_params({}), {}, )
def test_change_parameters_deny_invalid_params(self, load_module): workflow = Workflow.create_and_init() wf_module = workflow.tabs.first().wf_modules.create( order=0, slug="step-1", module_id_name="x", last_relevant_delta_id=workflow.last_delta_id, params={"x": 1}, ) ModuleVersion.create_or_replace_from_spec({ "id_name": "x", "name": "x", "category": "Clean", "parameters": [{ "id_name": "x", "type": "integer" }], }) load_module.return_value.param_schema = ParamDType.Dict( {"x": ParamDType.Integer()}) load_module.return_value.migrate_params = lambda x: x with self.assertRaises(ValueError): # Now the user requests to change params, giving an invalid param. self.run_with_async_db( commands.do( ChangeParametersCommand, workflow_id=workflow.id, wf_module=wf_module, new_values={"x": "Threeve"}, ))
def test_simple(self, load_module): load_module.return_value.migrate_params.return_value = {"A": "B"} load_module.return_value.fetch.return_value = FetchResult(self.output_path, []) result = fetch.fetch_or_wrap_error( self.ctx, self.chroot_context, self.basedir, WfModule(params={"A": "input"}, secrets={"C": "wrong"}), MockModuleVersion( id_name="A", param_schema=ParamDType.Dict({"A": ParamDType.String()}) ), {"C": "D"}, None, None, self.output_path, ) self.assertEqual(result, FetchResult(self.output_path, [])) load_module.return_value.migrate_params.assert_called_with({"A": "input"}) load_module.return_value.fetch.assert_called_with( chroot_context=self.chroot_context, basedir=self.basedir, params=Params({"A": "B"}), secrets={"C": "D"}, last_fetch_result=None, input_parquet_filename=None, output_filename=self.output_path.name, )
def param_schema(self) -> ParamDType.Dict: if "param_schema" in self.spec: # Module author wrote a schema in the YAML, to define storage of 'custom' parameters json_schema = self.spec["param_schema"] return ParamDType.parse({ "type": "dict", "properties": json_schema }) else: # Usual case: infer schema from module parameter types # Use of dict here means schema is not sensitive to parameter ordering, which is good return ParamDType.Dict( dict((f.id_name, f.dtype) for f in self.param_fields if f.dtype is not None))
def _load_tab_flows(workflow: Workflow, delta_id: int) -> List[TabFlow]: """ Query `workflow` for each tab's `TabFlow` (ordered by tab position). Raise `ModuleError` or `ValueError` if migrate_params() fails. Failed migration means the whole execute can't happen. """ ret = [] with workflow.cooperative_lock(): # reloads workflow if workflow.last_delta_id != delta_id: raise UnneededExecution for tab_model in workflow.live_tabs.all(): steps = [ ExecuteStep( step, (step.module_version.param_schema if step.module_version is not None else ParamDType.Dict({})), # We need to invoke the kernel and migrate _all_ modules' # params (WfModule.get_params), because we can only check # for tab cycles after migrating (and before calling any # render()). _get_migrated_params(step), ) for step in tab_model.live_wf_modules.all() ] ret.append(TabFlow(Tab(tab_model.slug, tab_model.name), steps)) return ret
def test_load_dynamic_ignore_test_py(self): code = b"def render(table, params):\n return table * 2" minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/imported.py", Body=code, ContentLength=len(code), ) # write other .py files that aren't module code and should be ignored minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/setup.py", Body=b"", ContentLength=0, ) minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/test_imported.py", Body=b"", ContentLength=0, ) with self.assertLogs("cjwstate.modules.loaded_module"): LoadedModule.for_module_version( MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now") )
def test_migrate_params_crash(self): code = b"def migrate_params(params):\n raise RuntimeError('xxx')" minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/imported.py", Body=code, ContentLength=len(code), ) with self.assertLogs("cjwstate.modules.loaded_module"): lm = LoadedModule.for_module_version( MockModuleVersion( "imported", "abcdef", ParamDType.Dict({"x": ParamDType.String()}), "now", ) ) with self.assertRaisesRegex(ModuleExitedError, "xxx"): # should have 'x' key lm.migrate_params({})
def test_load_dynamic_is_cached(self): code = b"def render(table, params):\n return table * 2" minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/imported.py", Body=code, ContentLength=len(code), ) with self.assertLogs("cjwstate.modules.loaded_module"): lm = LoadedModule.for_module_version( MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now") ) with patch("importlib.util.module_from_spec", None): lm2 = LoadedModule.for_module_version( MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now") ) self.assertIs(lm.compiled_module, lm2.compiled_module)
def test_load_dynamic(self): code = b"def render(table, params):\n return table * 2" minio.client.put_object( Bucket=minio.ExternalModulesBucket, Key="imported/abcdef/imported.py", Body=code, ContentLength=len(code), ) with self.assertLogs("cjwstate.modules.loaded_module"): lm = LoadedModule.for_module_version( MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now") ) self.assertEqual(lm.name, "imported:abcdef") # This ends up being kinda an integration test. with ExitStack() as ctx: basedir = Path(ctx.enter_context(tempdir_context(prefix="test-basedir-"))) basedir.chmod(0o755) input_table = ctx.enter_context( arrow_table_context({"A": [1]}, dir=basedir) ) input_table.path.chmod(0o644) output_tf = ctx.enter_context(tempfile.NamedTemporaryFile(dir=basedir)) ctx.enter_context(self.assertLogs("cjwstate.modules.loaded_module")) result = lm.render( basedir=basedir, input_table=input_table, params=Params({"col": "A"}), tab=Tab("tab-1", "Tab 1"), fetch_result=None, output_filename=Path(output_tf.name).name, ) assert_render_result_equals(result, RenderResult(arrow_table({"A": [2]})))
def test_change_parameters_across_module_versions(self, load_module): workflow = Workflow.create_and_init() # Initialize a WfModule that used module 'x' version '1' (which we # don't need to write in code -- after all, that version might be long # gone when ChangeParametersCommand is called. wf_module = workflow.tabs.first().wf_modules.create( order=0, slug="step-1", module_id_name="x", last_relevant_delta_id=workflow.last_delta_id, params={ "version": "v1", "x": 1 }, # version-'1' params ) # Now install version '2' of module 'x'. # # Version '2''s migrate_params() could do anything; in this test, it # simply changes 'version' from 'v1' to 'v2' ModuleVersion.create_or_replace_from_spec( { "id_name": "x", "name": "x", "category": "Clean", "parameters": [ { "id_name": "version", "type": "string" }, { "id_name": "x", "type": "integer" }, ], }, source_version_hash="2", ) load_module.return_value.param_dtype = ParamDType.Dict({ "version": ParamDType.String(), "x": ParamDType.Integer() }) load_module.return_value.migrate_params = lambda params: { **params, "version": "v2", } # Now the user requests to change params. # # The user was _viewing_ version '2' of module 'x', though # `wf_module.params` was at version 1. (Workbench ran # `migrate_params()` without saving the result when it # presented `params` to the user.) So the changes should apply atop # _migrated_ params. cmd = self.run_with_async_db( commands.do( ChangeParametersCommand, workflow_id=workflow.id, wf_module=wf_module, new_values={"x": 2}, )) self.assertEqual( wf_module.params, { "version": "v2", # migrate_params() ran "x": 2, # and we applied changes on top of its output }, ) self.run_with_async_db(commands.undo(cmd)) self.assertEqual( wf_module.params, { "version": "v1", "x": 1 } # exactly what we had before )
def test_load_static(self): # Test with a _real_ static module lm = LoadedModule.for_module_version( MockModuleVersion("pastecsv", "internal", ParamDType.Dict({}), "now") ) self.assertEqual(lm.name, "pastecsv:internal")