def test_generate_method_block(self): integration = Integration( service_name="SERVICE_NAME", class_prefix="CLASS_PREFIX", file_prefix="FILE_PREFFIX", client_path="CLOENT_PATH", ) method_block = hook_generator.generate_get_conn_method_block( integration) self.assertEqual( MethodBlock( name="get_conn", desc=[ "Retrieves client library object that allow access to SERVICE_NAME service." ], args={}, return_kind=TypeBrick(kind="CLOENT_PATH", indexes=[]), return_desc=["SERVICE_NAME client object."], code_blocks=[ CodeBlock( template_name="get_conn_body.py.tpl", template_params={ "client": TypeBrick(kind="CLOENT_PATH", indexes=[]) }, ) ], ), method_block, )
def test_should_parse_redis_update_instance(self, mock_docstring_parser): mock_docstring_parser.parse_docstring.return_value = ( self.parse_docstring_return_value) target_fn = CloudRedisClient.list_instances result = cloud_client_parser.parse_action_method( "update_instance", target_fn) mock_docstring_parser.parse_docstring.assert_called_once_with( "DOCSTRING") self.assertEqual( ActionInfo( name="update_instance", desc=["TEXT1"], args={ "parent": ParameterInfo( name="parent", kind=TypeBrick(kind="str", indexes=[]), desc=["TEXT2"], ), "instance_id": ParameterInfo( name="instance_id", kind=TypeBrick(kind="str", indexes=[]), desc=["TEXT3"], ), }, return_kind=None, return_desc=[ "A :class:`~google.cloud.redis_v1.types._OperationFuture` instance." ], ), result, )
def test_should_generate_value_for_optional_string( self, mock_generate_constant_list ): hook_method_blocks = [ MethodBlock( name="setUp", desc=None, args={ "arg3": ParameterBlock( "arg3", TypeBrick( kind="Optional", indexes=[TypeBrick(kind="str", indexes=[])] ), ) }, return_kind=None, return_desc=None, code_blocks=[], decorator_blocks=[], ) ] result = hook_test_generator.generate_constants(hook_method_blocks, []) mock_generate_constant_list.assert_called_once_with( {"arg3": TypeBrick(kind="Optional", indexes=[TypeBrick(kind="str")])} ) self.assertEqual("CONSTANTS", result)
def generate_ctor_method_block(client_info: ClientInfo, integration: Integration) -> MethodBlock: args = { "gcp_conn_id": ParameterBlock(name="gcp_conn_id", kind=TypeBrick("str")), "delegate_to": ParameterBlock(name="delegate_to", kind=TypeBrick("str")), } method_block = MethodBlock( name="__init__", desc=client_info.ctor_method.desc, args=args, return_kind=TypeBrick("None"), return_desc=None, code_blocks=[ CodeBlock( template_name="super_call.py.tpl", template_params={ "args": list(args.keys()), "kwargs": {} }, ), CodeBlock( template_name="set_field.py.tpl", template_params={ "field_name": "_client", "field_type": integration.client_type_brick, "field_value": "None", }, ), ], ) return method_block
def test_should_convert_with_not_matching_path_info(self): parameter = ParameterInfo( name="parent", kind=TypeBrick(kind="str", indexes=[]), desc=[ "Required. The resource name of the instance location using the form: " "``projects/{project_id}/locations/{location_id}`` where " "``locati`` refers to a GCP region" ], ) path_infos = { "location": PathInfo(name="location", args=["not_project", "not_location"]) } integration = Integration( service_name="SERVICE_NAME", class_prefix="CLASS_PREFIX", file_prefix="FILE_PREFFIX", client_path="CLOENT_PATH", ) req_params, opt_params, code = hook_generator.convert_path_parameter_block_to_individual_parameters( parameter, path_infos=path_infos, integration=integration) self.assertEqual( { "location": ParameterBlock( name="location", kind=TypeBrick(kind="str", indexes=[]), desc=["TODO: Fill description"], default_value=None, ) }, req_params, ) self.assertEqual( { "project_id": ParameterBlock( name="project_id", kind=TypeBrick(kind="str", indexes=[]), desc=["TODO: Fill description"], default_value="None", ) }, opt_params, ) self.assertEqual( CodeBlock( template_name="call_path.py.tpl", template_params={ "var_name": "parent", "fn_name": "TODO", "args": ["project_id", "location"], "client": TypeBrick(kind="CLOENT_PATH", indexes=[]), }, ), code, )
def test_should_parse_dict(self): self.assertEqual( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ), typehint_parser.parse_typehint("Dict[str, str]"), )
def test_should_fix_incorrect_type(self): self.assertEqual( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="float", indexes=[]), ], ), typehint_parser.parse_typehint("dict[str -> float]"), )
def test_should_refinement_nested_dict(self): self.assertEqual( TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Optional", indexes=[TypeBrick(kind="Dict", indexes=[])]) ], ), typehint_parser.parse_typehint("Union[Optional[dict]]"), )
def test_generate_method_block(self): client_info = mock.MagicMock( **{"ctor_method.desc": ["DESC_A", "DESC_B"]}) integration = Integration( service_name="SERVICE_NAME", class_prefix="CLASS_PREFIX", file_prefix="FILE_PREFFIX", client_path="CLOENT_PATH", ) method_block = hook_generator.generate_ctor_method_block( client_info, integration) self.assertEqual( MethodBlock( name="__init__", desc=["DESC_A", "DESC_B"], args={ "gcp_conn_id": ParameterBlock( name="gcp_conn_id", kind=TypeBrick(kind="str", indexes=[]), desc=None, default_value=None, ), "delegate_to": ParameterBlock( name="delegate_to", kind=TypeBrick(kind="str", indexes=[]), desc=None, default_value=None, ), }, return_kind=TypeBrick(kind="None", indexes=[]), return_desc=None, code_blocks=[ CodeBlock( template_name="super_call.py.tpl", template_params={ "args": ["gcp_conn_id", "delegate_to"], "kwargs": {}, }, ), CodeBlock( template_name="set_field.py.tpl", template_params={ "field_name": "_client", "field_type": TypeBrick(kind="CLOENT_PATH", indexes=[]), "field_value": "None", }, ), ], ), method_block, )
def test_create_constants(self, mock_generate_constant_list): class_a = ClassBlock( name="CLASS", extend_class="EXTEND_CLASS", methods_blocks=[ MethodBlock( name="__init__", desc=None, args={"ARG_A": ParameterBlock("ARG_A", kind=TypeBrick("str"))}, return_kind=None, return_desc=None, code_blocks=[], decorator_blocks=[], ) ], ) class_b = ClassBlock( name="CLASS", extend_class="EXTEND_CLASS", methods_blocks=[ MethodBlock( name="__init__", desc=None, args={ "ARG_B": ParameterBlock("ARG_B"), "*args": ParameterBlock("*args"), "**kwargs": ParameterBlock("**kwargs"), }, return_kind=None, return_desc=None, code_blocks=[], decorator_blocks=[], ), MethodBlock( name="execute", desc=None, args={"ARG_C": ParameterBlock("ARG_C")}, return_kind=None, return_desc=None, code_blocks=[], decorator_blocks=[], ), ], ) result = operator_test_generator.create_constants([class_a, class_b]) self.assertEqual("CONSTANTS", result) mock_generate_constant_list.assert_called_once_with( {"ARG_A": TypeBrick(kind="str"), "ARG_B": None} )
def _create_action_info(prefix: str): return ActionInfo( name=f"{prefix}NAME", desc=[f"{prefix}DESC_A", f"{prefix}DESC_B"], args={ f"{prefix}ARG_A": ParameterInfo( name=f"{prefix}ARG_A", kind=TypeBrick("str"), desc=[f"{prefix}DESC_C", f"{prefix}DESC_D"], ) }, return_kind=TypeBrick("float"), return_desc=[f"{prefix}DESC_E", f"{prefix}DESC_F"], )
def generate_test_method_for_call_without_project_id( hook_class_path: str, hook_method_block: MethodBlock) -> MethodBlock: code_blocks = [ CodeBlock( template_name="method_call_assert_raises.py.tpl", template_params={ "target": f"self.hook.{hook_method_block.name}", "call_params": { a: f"TEST_{a.upper()}" if a != "project_id" else "None" for a in hook_method_block.args.keys() }, }, ) ] method_block = MethodBlock( name=f"test_{hook_method_block.name}_without_project_id", desc=None, args={"mock_get_conn": ParameterBlock("mock_get_conn")}, return_kind=TypeBrick("None"), return_desc=None, code_blocks=code_blocks, decorator_blocks=[ CodeBlock( template_name="decorator_mock_get_conn.py.tpl", template_params={"class_path": hook_class_path}, ) ], ) return method_block
def test_should_parse_deprecated_union_style(self): self.assertEqual( TypeBrick( kind="Iterator", indexes=[ TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Dict", indexes=[]), TypeBrick(kind="Set", indexes=[]), ], ) ], ), typehint_parser.parse_typehint("iterator[dict|set]"), )
def _read_type_element(self) -> TypeBrick: assert self.hint_iter is not None kind = self._read_kind() indexes: List[TypeBrick] = [] if self.hint_iter.has_next() and self.hint_iter.peek() == "[": indexes = self._read_type_list() return TypeBrick(kind, indexes)
def create_assert_call_method_block( operator_module_path: str, hook_class_name: str, operator_class_block: ClassBlock ) -> MethodBlock: ctor_method = next( m for m in operator_class_block.methods_blocks if m.name == "__init__" ) ctor_args = [a for a in ctor_method.args.keys() if a not in ("*args", "**kwargs")] method_name = find_client_call_method_name(operator_class_block) return MethodBlock( name="test_assert_valid_hook_call", desc=None, args={"mock_hook": ParameterBlock("mock_hook")}, return_kind=TypeBrick("None"), return_desc=None, code_blocks=[ CodeBlock( template_name="method_call.py.tpl", template_params={ "var_name": "task", "target": operator_class_block.name, "call_params": {a: f"TEST_{a.upper()}" for a in ctor_args}, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "task.execute", "call_params": {"context": "mock.MagicMock()"}, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "mock_hook.assert_called_once_with", "call_params": {"gcp_conn_id": "TEST_GCP_CONN_ID"}, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": f"mock_hook.return_value.{method_name}.assert_called_once_with", "call_params": { a: f"TEST_{a.upper()}" for a in ctor_args if a != "gcp_conn_id" }, }, ), ], decorator_blocks=[ CodeBlock( "decorator_mock_hook.py.tpl", template_params={ "class_path": f"{operator_module_path}.{hook_class_name}" }, ) ], )
def walk_recursive_and_replace(element: TypeBrick): if not element.indexes and "|" in element.kind: element.indexes = [ TypeBrick(kind=index.strip()) for index in element.kind.split("|") ] element.kind = "Union" if element.indexes: for child in element.indexes: walk_recursive_and_replace(child)
def convert_path_parameter_block_to_individual_parameters( path_parameter: ParameterInfo, path_infos: Dict[str, PathInfo], integration: Integration, ) -> Tuple[Dict[str, ParameterBlock], Dict[str, ParameterBlock], CodeBlock]: match = PATH_REGEXP.search("\n".join(path_parameter.desc)) if not match: raise GeneratorException("") path = match["path"] path_info = find_matching_path_info(path, path_infos=path_infos) optional_parameters: Dict[str, ParameterBlock] = {} required_parameters: Dict[str, ParameterBlock] = {} call_args = [] for name in path_info.args: if name == "project": optional_parameters["project_id"] = ParameterBlock( name="project_id", kind=TypeBrick("str"), desc=["TODO: Fill description"], default_value="None", ) call_args.append("project_id") else: required_parameters[name] = ParameterBlock( name=name, kind=TypeBrick("str"), desc=["TODO: Fill description"]) call_args.append(name) code_block = CodeBlock( template_name="call_path.py.tpl", template_params={ "var_name": path_parameter.name, "fn_name": path_info.name, "args": call_args, "client": integration.client_type_brick, }, ) return required_parameters, optional_parameters, code_block
def test_generate_test_method_for_client_call(self): action_info = self._create_action("ACTION", ["ARG_1", "ARG_2"]) method_block = self._create_method("METHOD", ["ARG_2", "ARG_3"]) hook_class_path = "HOOK_CLASS_PATH" result = hook_test_generator.generate_test_method_for_client_call( action_info, method_block, hook_class_path, "PROJECT_ID_VALUE" ) self.assertEqual( MethodBlock( name="test_METHOD", desc=None, args={ "mock_get_conn": ParameterBlock( name="mock_get_conn", kind=None, desc=None, default_value=None ) }, return_kind=TypeBrick(kind="None", indexes=[]), return_desc=None, code_blocks=[ CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "self.hook.METHOD", "call_params": { "ARG_2": "TEST_ARG_2", "ARG_3": "TEST_ARG_3", }, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "mock_get_conn.ACTION.assert_called_once_with", "call_params": { "ARG_1": "TEST_ARG_1", "ARG_2": "TEST_ARG_2", }, }, ), ], decorator_blocks=[ CodeBlock( template_name="decorator_mock_get_conn.py.tpl", template_params={"class_path": "HOOK_CLASS_PATH"}, ) ], ), result, )
def generate_test_method_for_client_call( action_info: ActionInfo, hook_method_block: MethodBlock, hook_class_path: str, project_id_value: str, ): code_blocks = [ CodeBlock( template_name="method_call.py.tpl", template_params={ "target": f"self.hook.{hook_method_block.name}", "call_params": { arg_name: f"TEST_{arg_name.upper()}" for arg_name in hook_method_block.args.keys() }, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": f"mock_get_conn.{action_info.name}.assert_called_once_with", "call_params": { arg_name: f"TEST_{arg_name.upper()}" if arg_name != "project_id" else project_id_value for arg_name in action_info.args.keys() }, }, ), ] method_block = MethodBlock( name=f"test_{hook_method_block.name}", desc=None, args={"mock_get_conn": ParameterBlock("mock_get_conn")}, return_kind=TypeBrick("None"), return_desc=None, code_blocks=code_blocks, decorator_blocks=[ CodeBlock( template_name="decorator_mock_get_conn.py.tpl", template_params={"class_path": hook_class_path}, ) ], ) return method_block
def generate_setup_method_block(class_path: str, init_new: str) -> MethodBlock: method_block = MethodBlock( name="setUp", desc=None, args={}, return_kind=TypeBrick("None"), return_desc=None, code_blocks=[ CodeBlock( template_name="setup_mock.py.tpl", template_params={ "class_path": class_path, "new": init_new }, ) ], ) return method_block
def generate_constant_list( unique_constant: Dict[str, Optional[TypeBrick]] ) -> List[Constant]: constants = [] for name, kind in unique_constant.items(): constant_name = f"TEST_{name.upper()}" constant_kind = kind or TypeBrick("None") if constant_kind and (constant_kind.is_union or constant_kind.is_optional): constant_kind = constant_kind.indexes[0] constant_value = ( f'\'test-{name.replace("_", "-")}\'' if constant_kind.name == "str" else "None # TODO: Fill missing value" ) constants.append( Constant(name=constant_name, kind=constant_kind, value=constant_value) ) return constants
def test_generate_setup_method_block(self): result = hook_test_generator.generate_setup_method_block( class_path="CLASS_PATH", init_new="INIT_NEW" ) self.assertEqual( MethodBlock( name="setUp", desc=None, args={}, return_kind=TypeBrick(kind="None", indexes=[]), return_desc=None, code_blocks=[ CodeBlock( template_name="setup_mock.py.tpl", template_params={"class_path": "CLASS_PATH", "new": "INIT_NEW"}, ) ], decorator_blocks=[], ), result, )
def test_generate_test_method_for_call_without_project_id(self): method_a = self._create_method("TEST_METHOD", ["ARG_A", "project_id"]) result = hook_test_generator.generate_test_method_for_call_without_project_id( "HOOK_CLASS_PATH", method_a ) self.assertEqual( MethodBlock( name="test_TEST_METHOD_without_project_id", desc=None, args={ "mock_get_conn": ParameterBlock( name="mock_get_conn", kind=None, desc=None, default_value=None ) }, return_kind=TypeBrick(kind="None", indexes=[]), return_desc=None, code_blocks=[ CodeBlock( template_name="method_call_assert_raises.py.tpl", template_params={ "target": "self.hook.TEST_METHOD", "call_params": { "ARG_A": "TEST_ARG_A", "project_id": "None", }, }, ) ], decorator_blocks=[ CodeBlock( template_name="decorator_mock_get_conn.py.tpl", template_params={"class_path": "HOOK_CLASS_PATH"}, ) ], ), result, )
class TestTypeHintParser(unittest.TestCase): def test_should_parse_primitives(self): self.assertEqual(TypeBrick(kind="str", indexes=[]), typehint_parser.parse_typehint("str")) def test_should_parse_class(self): self.assertEqual( TypeBrick(kind="airflow_munchkin.typehint_parser", indexes=[]), typehint_parser.parse_typehint("airflow_munchkin.typehint_parser"), ) def test_should_parse_list(self): self.assertEqual( TypeBrick(kind="List", indexes=[TypeBrick(kind="str", indexes=[])]), typehint_parser.parse_typehint("List[str]"), ) def test_should_parse_dict(self): self.assertEqual( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ), typehint_parser.parse_typehint("Dict[str, str]"), ) def test_should_fix_incorrect_type(self): self.assertEqual( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="float", indexes=[]), ], ), typehint_parser.parse_typehint("dict[str -> float]"), ) def test_should_refinement_dict(self): self.assertEqual(TypeBrick(kind="Dict", indexes=[]), typehint_parser.parse_typehint("dict")) def test_should_refinement_nested_dict(self): self.assertEqual( TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Optional", indexes=[TypeBrick(kind="Dict", indexes=[])]) ], ), typehint_parser.parse_typehint("Union[Optional[dict]]"), ) def test_should_parse_deprecated_union_style(self): self.assertEqual( TypeBrick( kind="Iterator", indexes=[ TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Dict", indexes=[]), TypeBrick(kind="Set", indexes=[]), ], ) ], ), typehint_parser.parse_typehint("iterator[dict|set]"), ) @parameterized.expand([ (TypeBrick(kind="str", indexes=[]), "str"), ( TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Dict", indexes=[]), TypeBrick( kind= "google.cloud.automl_v1beta1.types.BatchPredictInputConfig", indexes=[], ), ], ), "Union[dict, ~google.cloud.automl_v1beta1.types.BatchPredictInputConfig]", ), ( TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Dict", indexes=[]), TypeBrick( kind= "google.cloud.automl_v1beta1.types.BatchPredictOutputConfig", indexes=[], ), ], ), "Union[dict, ~google.cloud.automl_v1beta1.types.BatchPredictOutputConfig]", ), ( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ), "Dict[str, str]", ), ( TypeBrick( kind="Optional", indexes=[ TypeBrick(kind="google.api_core.retry.Retry", indexes=[]) ], ), "Optional[google.api_core.retry.Retry]", ), ( TypeBrick(kind="Optional", indexes=[TypeBrick(kind="float", indexes=[])]), "Optional[float]", ), ( TypeBrick( kind="Optional", indexes=[ TypeBrick( kind="Sequence", indexes=[ TypeBrick( kind="Tuple", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ) ], ) ], ), "Optional[Sequence[Tuple[str, str]]]", ), (TypeBrick(kind="str", indexes=[]), "str"), ( TypeBrick( kind="Union", indexes=[ TypeBrick(kind="Dict", indexes=[]), TypeBrick( kind="google.cloud.automl_v1beta1.types.ExamplePayload", indexes=[], ), ], ), "Union[dict, ~google.cloud.automl_v1beta1.types.ExamplePayload]", ), ( TypeBrick( kind="Dict", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ), "Dict[str, str]", ), ( TypeBrick( kind="Optional", indexes=[ TypeBrick(kind="google.api_core.retry.Retry", indexes=[]) ], ), "Optional[google.api_core.retry.Retry]", ), ( TypeBrick( kind="Optional", indexes=[ TypeBrick( kind="Sequence", indexes=[ TypeBrick( kind="Tuple", indexes=[ TypeBrick(kind="str", indexes=[]), TypeBrick(kind="str", indexes=[]), ], ) ], ) ], ), "Optional[Sequence[Tuple[str, str]]]", ), ]) def test_real_cases(self, exptected_result, input_text): self.assertEqual(exptected_result, typehint_parser.parse_typehint(input_text))
def test_should_refinement_dict(self): self.assertEqual(TypeBrick(kind="Dict", indexes=[]), typehint_parser.parse_typehint("dict"))
def test_should_parse_list(self): self.assertEqual( TypeBrick(kind="List", indexes=[TypeBrick(kind="str", indexes=[])]), typehint_parser.parse_typehint("List[str]"), )
def test_should_parse_class(self): self.assertEqual( TypeBrick(kind="airflow_munchkin.typehint_parser", indexes=[]), typehint_parser.parse_typehint("airflow_munchkin.typehint_parser"), )
def test_should_parse_primitives(self): self.assertEqual(TypeBrick(kind="str", indexes=[]), typehint_parser.parse_typehint("str"))
def test_create_operator_class_block(self): hook_class_name = "HOOK_CLASS_NAME" integration = Integration( service_name="SERVICE_NAME", class_prefix="CLASS_PREFIX", file_prefix="FILE_PREFFIX", client_path="CLOENT_PATH", ) hook_method = self._create_method("METHOD_A", ["arg_a", "arg_b"]) result = operator_generator.create_operator_class_block( hook_class_name, hook_method, integration=integration) self.assertEqual( ClassBlock( name="CLASS_PREFIXMethodAOperator", extend_class="airflow.models.BaseOperator", methods_blocks=[ MethodBlock( name="__init__", desc=None, args={ "arg_a": ParameterBlock(name="arg_a", kind=None, desc=None, default_value=None), "arg_b": ParameterBlock(name="arg_b", kind=None, desc=None, default_value=None), "gcp_conn_id": ParameterBlock( name="gcp_conn_id", kind=TypeBrick(kind="str"), desc=None, default_value="'google_cloud_default'", ), "*args": ParameterBlock(name="*args", kind=None, desc=None, default_value=None), "**kwargs": ParameterBlock( name="**kwargs", kind=None, desc=None, default_value=None, ), }, return_kind=TypeBrick(kind="None"), return_desc=None, code_blocks=[ CodeBlock( template_name="super_call.py.tpl", template_params={ "args": ["*args", "**kwargs"], "kwargs": {}, }, ), CodeBlock( template_name="set_field.py.tpl", template_params={ "field_name": "arg_a", "field_value": "arg_a", "field_type": None, }, ), CodeBlock( template_name="set_field.py.tpl", template_params={ "field_name": "arg_b", "field_value": "arg_b", "field_type": None, }, ), CodeBlock( template_name="set_field.py.tpl", template_params={ "field_name": "gcp_conn_id", "field_value": "gcp_conn_id", "field_type": None, }, ), ], decorator_blocks=[ CodeBlock( template_name="decorator_apply_defaults.py.tpl", template_params={}, ) ], ), MethodBlock( name="execute", desc=None, args={ "context": ParameterBlock( name="context", kind=TypeBrick(kind="Dict"), desc=None, default_value=None, ) }, return_kind=None, return_desc=None, code_blocks=[ CodeBlock( template_name="method_call.py.tpl", template_params={ "var_name": "hook", "target": "HOOK_CLASS_NAME", "call_params": { "gcp_conn_id": "self.gcp_conn_id" }, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "hook.METHOD_A", "call_params": { "arg_a": "self.arg_a", "arg_b": "self.arg_b", }, }, ), ], decorator_blocks=[], ), ], ), result, )
def test_generate_method_block( self, mock_convert_path_parameter_block_to_individual_parameters): action_info = ActionInfo( name="NAME", desc=["DESC_A", "DESC_B"], args={ "ARG_A": ParameterInfo(name="ARG_A", kind=TypeBrick("str"), desc=["DESC_C", "DESC_D"]) }, return_kind=TypeBrick("float"), return_desc=["DESC_E", "DESC_F"], ) integration = Integration( service_name="SERVICE_NAME", class_prefix="CLASS_PREFIX", file_prefix="FILE_PREFFIX", client_path="CLOENT_PATH", ) method_block = hook_generator.generate_method_block( action_info, path_infos="PATH_INFOS", integration=integration) self.assertEqual( MethodBlock( name="NAME", desc=["DESC_A", "DESC_B"], args={ "ARG_A": ParameterBlock( name="ARG_A", kind=TypeBrick(kind="str", indexes=[]), desc=["DESC_C", "DESC_D"], default_value=None, ) }, return_kind=TypeBrick(kind="float", indexes=[]), return_desc=["DESC_E", "DESC_F"], code_blocks=[ CodeBlock( template_name="method_call.py.tpl", template_params={ "var_name": "client", "target": "self.get_conn", "call_params": {}, }, ), CodeBlock( template_name="method_call.py.tpl", template_params={ "target": "client.NAME", "var_name": "result", "call_params": { "ARG_A": "ARG_A" }, }, ), CodeBlock( template_name="return.py.tpl", template_params={"var_name": "result"}, ), ], decorator_blocks=[], ), method_block, )