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