def test_find_client_call_method_name(self):
operator_class_block = ClassBlock(
name="CloudMemorystoreCreateInstanceOperator",
extend_class="airflow.models.BaseOperator",
methods_blocks=[
MethodBlock(
name="execute",
desc=None,
args={},
return_kind=None,
return_desc=None,
code_blocks=[
CodeBlock(
template_name="method_call.py.tpl",
template_params={"target": "CloudMemorystoreHook"},
),
CodeBlock(
template_name="method_call.py.tpl",
template_params={
"target": "hook.create_instance",
"call_params": {},
},
),
],
decorator_blocks=[],
)
],
)
result = operator_test_generator.find_client_call_method_name(
operator_class_block
)
self.assertEqual("create_instance", result)
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 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 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 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_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 test_render_code_block(self, mock_render_template):
code_block = CodeBlock(template_name="TEMPLATE_NAME",
template_params={"PARAM_A": "VALUE_A"})
result = brushes.render_code_block(code_block)
mock_render_template.assert_called_once_with(
PARAM_A="VALUE_A", template_name="code_blocks/TEMPLATE_NAME")
self.assertEqual("TEMPLATE", result)
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_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 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 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,
)
def generate_get_conn_method_block(integration: Integration) -> MethodBlock:
method_block = MethodBlock(
name="get_conn",
desc=[
f"Retrieves client library object that allow access to {integration.service_name} service."
],
args={},
return_kind=integration.client_type_brick,
return_desc=[f"{integration.service_name} client object."],
code_blocks=[
CodeBlock(
template_name="get_conn_body.py.tpl",
template_params={"client": integration.client_type_brick},
)
],
)
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 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 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_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,
)
class TestGenerateMethodBlock(TestCase):
@mock.patch(
"airflow_munchkin.block_generator.hook_generator.convert_path_"
"parameter_block_to_individual_parameters",
return_value=(
{
"location":
ParameterBlock(
name="location",
kind=TypeBrick(kind="str", indexes=[]),
desc=["TODO: Fill description"],
default_value=None,
)
},
{
"project_id":
ParameterBlock(
name="project_id",
kind=TypeBrick(kind="str", indexes=[]),
desc=["TODO: Fill description"],
default_value="None",
)
},
CodeBlock(
template_name="call_path.py.tpl",
template_params={
"var_name": "parent",
"fn_name": "location",
"args": ["project_id", "location"],
"client": TypeBrick(kind="CLOENT_PATH", indexes=[]),
},
),
),
)
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,
)
def test_create_assert_call_method_block(self, mock_find_client_call_method_name):
operator_module_path = "OPERATOR_MODULE_PATH"
hook_class_name = "CloudMemorystoreHook"
operator_class_block = ClassBlock(
name="CloudMemorystoreCreateInstanceOperator",
extend_class="airflow.models.BaseOperator",
methods_blocks=[
MethodBlock(
name="__init__",
desc=[
"Creates a Redis instance based on the specified tier and memory size."
],
args={
"location": ParameterBlock(
name="location",
kind=TypeBrick(kind="str"),
desc=["TODO: Fill description"],
default_value=None,
),
"instance_id": ParameterBlock(
name="instance_id", kind=TypeBrick(kind="str")
),
"instance": ParameterBlock(
name="instance",
kind=TypeBrick(
kind="Union",
indexes=[
TypeBrick(kind="Dict"),
TypeBrick(
kind="google.cloud.redis_v1.types.Instance"
),
],
),
),
"project_id": ParameterBlock(
name="project_id",
kind=TypeBrick(kind="str"),
desc=["TODO: Fill description"],
default_value="None",
),
"retry": ParameterBlock(
name="retry",
kind=TypeBrick(kind="google.api_core.retry.Retry"),
default_value="None",
),
"timeout": ParameterBlock(
name="timeout",
kind=TypeBrick(kind="float"),
default_value="None",
),
"metadata": ParameterBlock(
name="metadata",
kind=TypeBrick(
kind="Sequence",
indexes=[
TypeBrick(
kind="Tuple",
indexes=[
TypeBrick(kind="str"),
TypeBrick(kind="str"),
],
)
],
),
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"),
"**kwargs": ParameterBlock(name="**kwargs"),
},
return_kind=TypeBrick(kind="None"),
return_desc=None,
code_blocks=[],
decorator_blocks=[],
)
],
)
result = operator_test_generator.create_assert_call_method_block(
operator_module_path, hook_class_name, operator_class_block
)
self.assertEqual(
MethodBlock(
name="test_assert_valid_hook_call",
desc=None,
args={"mock_hook": ParameterBlock(name="mock_hook")},
return_kind=TypeBrick(kind="None"),
return_desc=None,
code_blocks=[
CodeBlock(
template_name="method_call.py.tpl",
template_params={
"var_name": "task",
"target": "CloudMemorystoreCreateInstanceOperator",
"call_params": {
"location": "TEST_LOCATION",
"instance_id": "TEST_INSTANCE_ID",
"instance": "TEST_INSTANCE",
"project_id": "TEST_PROJECT_ID",
"retry": "TEST_RETRY",
"timeout": "TEST_TIMEOUT",
"metadata": "TEST_METADATA",
"gcp_conn_id": "TEST_GCP_CONN_ID",
},
},
),
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": "mock_hook.return_value.METHOD_NAME.assert_called_once_with",
"call_params": {
"location": "TEST_LOCATION",
"instance_id": "TEST_INSTANCE_ID",
"instance": "TEST_INSTANCE",
"project_id": "TEST_PROJECT_ID",
"retry": "TEST_RETRY",
"timeout": "TEST_TIMEOUT",
"metadata": "TEST_METADATA",
},
},
),
],
decorator_blocks=[
CodeBlock(
template_name="decorator_mock_hook.py.tpl",
template_params={
"class_path": "OPERATOR_MODULE_PATH.CloudMemorystoreHook"
},
)
],
),
result,
)
mock_find_client_call_method_name.assert_called_once_with(operator_class_block)
def create_operator_class_block(hook_class_name: str, hook_method: MethodBlock,
integration: Integration) -> ClassBlock:
init_args = {
**dict(hook_method.args.items()),
"gcp_conn_id":
ParameterBlock("gcp_conn_id",
kind=TypeBrick("str"),
default_value="'google_cloud_default'"),
}
super_args = {
"*args": ParameterBlock("*args"),
"**kwargs": ParameterBlock("**kwargs"),
}
ctor_method_block = MethodBlock(
name="__init__",
desc=hook_method.desc,
args={
**init_args,
**super_args
},
return_kind=TypeBrick("None"),
return_desc=None,
code_blocks=[
CodeBlock(
template_name="super_call.py.tpl",
template_params={
"args": list(super_args.keys()),
"kwargs": {}
},
),
*(CodeBlock(
template_name="set_field.py.tpl",
template_params={
"field_name": arg_name,
"field_value": arg_name,
"field_type": None,
},
) for arg_name in init_args.keys()),
],
decorator_blocks=[
CodeBlock(template_name="decorator_apply_defaults.py.tpl",
template_params={})
],
)
execute_method_block = MethodBlock(
name="execute",
desc=None,
args={"context": ParameterBlock("context", TypeBrick("Dict"))},
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": f"hook.{hook_method.name}",
"call_params": {
arg_name: f"self.{arg_name}"
for arg_name in hook_method.args
},
},
),
],
decorator_blocks=[
CodeBlock(template_name="decorator_apply_defaults.py.tpl",
template_params={})
],
)
return ClassBlock(
name=
f"{integration.class_prefix}{utils.to_camel_case(hook_method.name)}Operator",
extend_class="airflow.models.BaseOperator",
methods_blocks=[ctor_method_block, execute_method_block],
)
def generate_method_block(action_info: ActionInfo, path_infos: Dict[str,
PathInfo],
integration: Integration) -> MethodBlock:
blocks: List[CodeBlock] = []
blocks.append(
CodeBlock(
template_name="method_call.py.tpl",
template_params=dict(var_name="client",
target="self.get_conn",
call_params={}),
))
optional_parameters: Dict[str, ParameterBlock] = {}
required_parameters: Dict[str, ParameterBlock] = {}
for arg in action_info.args.values():
if arg.name == "parent" or arg.name == "name":
new_req_params, new_opt_params, code = convert_path_parameter_block_to_individual_parameters(
path_parameter=arg,
path_infos=path_infos,
integration=integration)
required_parameters.update(new_req_params)
optional_parameters.update(new_opt_params)
blocks.append(code)
continue
kind = arg.kind
default_value = None
if kind and kind.is_optional:
kind = kind.indexes[0]
default_value = "None"
optional_parameters[arg.name] = ParameterBlock(
name=arg.name,
kind=kind,
desc=arg.desc,
default_value=default_value)
else:
required_parameters[arg.name] = ParameterBlock(
name=arg.name,
kind=kind,
desc=arg.desc,
default_value=default_value)
method_block = MethodBlock(
name=action_info.name,
desc=action_info.desc,
args={
**required_parameters,
**optional_parameters
},
return_kind=action_info.return_kind,
return_desc=action_info.return_desc,
code_blocks=blocks,
)
blocks.append(
CodeBlock(
template_name="method_call.py.tpl",
template_params=dict(
target=f"client.{action_info.name}",
var_name="result" if action_info.return_desc else None,
call_params={arg: arg
for arg in action_info.args},
),
))
if action_info.return_desc:
blocks.append(
CodeBlock(
template_name="return.py.tpl",
template_params=dict(
var_name="result" if action_info.return_desc else None),
))
return method_block