def test_generate_class_block_without_default_project_id(
self,
mock_generate_setup_method_block,
mock_generate_test_method_for_client_call,
mock_generate_test_method_for_call_without_project_id,
):
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLIENT_PATH",
)
method_a = self._create_method("METHOD_A", args=["location"])
method_b = self._create_method("METHOD_B", args=["project_id"])
action_a = self._create_action("METHOD_A")
action_b = self._create_action("METHOD_B")
result = hook_test_generator.generate_class_block_without_default_project_id(
hook_class_path="HOOK_CLASS_PATH",
hook_class=ClassBlock(
name="CLASS_PREFIXHook",
extend_class="airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
methods_blocks=[method_a, method_b],
),
client_info=ClientInfo(
ctor_method=None,
path_methods={},
action_methods={"METHOD_A": action_a, "METHOD_B": action_b},
),
integration=integration,
)
self.assertEqual(
ClassBlock(
name="TestCLASS_PREFIXWithoutDefaultProjectIdHook",
extend_class="unittest.TestCase",
methods_blocks=["CLASS_A", "TEST_A", "TEST_B", "TEST_C"],
),
result,
)
mock_generate_setup_method_block.assert_called_once_with(
"HOOK_CLASS_PATH", "mock_base_gcp_hook_no_default_project_id"
)
mock_generate_test_method_for_client_call.assert_any_call(
action_a, method_a, "HOOK_CLASS_PATH", "TEST_PROJECT_ID"
)
mock_generate_test_method_for_client_call.assert_any_call(
action_b, method_b, "HOOK_CLASS_PATH", "TEST_PROJECT_ID"
)
mock_generate_test_method_for_call_without_project_id.assert_any_call(
"HOOK_CLASS_PATH", method_b
)
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 test_render_class_block(self, mock_render_template, mock_render_class_block):
ctor_method = mock.MagicMock()
ctor_method.name = "__init__"
ctor_method.args = {"*args": "ParameterBlock", "**kwargs": "ParameterBlock"}
method_a = mock.MagicMock()
method_a.name = "method_a"
method_b = mock.MagicMock()
method_b.name = "method_b"
class_block = ClassBlock(
name="NAME",
extend_class="EXTEND_CLASS",
methods_blocks=[ctor_method, method_a, method_b],
)
result = brushes.render_class_block(class_block)
mock_render_template.assert_called_once_with(
ctor_method=ctor_method,
ctor_method_args_docstring=[],
extend_class="EXTEND_CLASS",
method_blocks=[
"TEMPLATE_METHOD_BLOCK",
"TEMPLATE_METHOD_BLOCK",
"TEMPLATE_METHOD_BLOCK",
],
name="NAME",
template_name="class_block.py.tpl",
)
mock_render_class_block.assert_any_call(method_a)
mock_render_class_block.assert_any_call(method_b)
self.assertEqual("TEMPLATE", result)
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_class_block_with_default_project_id(
hook_class_path: str,
hook_class: ClassBlock,
client_info: ClientInfo,
integration: Integration,
) -> ClassBlock:
hook_methods = []
hook_methods.append(
generate_setup_method_block(hook_class_path,
"mock_base_gcp_hook_default_project_id"))
for name in client_info.action_methods.keys():
action_info = client_info.action_methods[name]
hook_method_block = next(m for m in hook_class.methods_blocks
if m.name == name)
hook_methods.append(
generate_test_method_for_client_call(
action_info,
hook_method_block,
hook_class_path,
"GCP_PROJECT_ID_HOOK_UNIT_TEST",
))
class_block = ClassBlock(
name=f"Test{integration.class_prefix}WithDefaultProjectIdHook",
extend_class="unittest.TestCase",
methods_blocks=hook_methods,
)
return class_block
def test_visit_class_should_remember_imports(self):
gather = ImportGather()
class_block = ClassBlock(
name="NAME", extend_class="awesome.cat", methods_blocks=[]
)
gather.visit_class(class_block)
self.assertEqual({"awesome.cat"}, gather.import_statement)
def create_operator_test_class_blocks(
operator_module_path: str, hook_class_name: str, operator_class: ClassBlock
) -> ClassBlock:
return ClassBlock(
name=f"Test{operator_class.name}",
extend_class="unittest.TestCase",
methods_blocks=[
create_assert_call_method_block(
operator_module_path, hook_class_name, operator_class
)
],
)
def test_create_file_block(self, mock_create_operator_class_blocks,
mock_imports_statement_gather):
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLOENT_PATH",
)
hook_class_block = ClassBlock(
name="CLASS_PREFIXHook",
extend_class=
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
methods_blocks=["METHOD_CTOR", "METHOD_GET_CONN", "METHOD_A"],
)
hook_file_block = FileBlock(
file_name="FILE_PREFFIX_hook.py",
class_blocks=[hook_class_block],
import_statement={"IMPORT_A"},
)
result = operator_generator.create_file_block(integration,
hook_file_block)
self.assertEqual(
FileBlock(
file_name="FILE_PREFFIX_operator.py",
class_blocks=["CLASS_A", "CLASS_B"],
import_statement={
"typing.Tuple",
"typing.Union",
"typing.Optional",
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
"unittest.mock",
"google.cloud.redis_v1beta1.CloudRedisClient",
"airflow.utils.decorators.apply_defaults",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_default_project_id",
"airflow.contrib.hooks.FILE_PREFFIX_hook.CLASS_PREFIXHook",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_no_default_project_id",
"typing.Dict",
"tests.contrib.utils.base_gcp_mock.GCP_PROJECT_ID_HOOK_UNIT_TEST",
"airflow.AirflowException",
"google.api_core.retry.Retry",
"typing.Sequence",
"unittest.TestCase",
},
constants=[],
),
result,
)
mock_imports_statement_gather.update_imports_statements.assert_called_once_with(
result)
mock_create_operator_class_blocks.assert_called_once_with(
hook_class_block, integration)
def test_create_operator_test_class_blocks(mock_create_assert_call_method_block):
operator_module_path = "OPERATOR_MODULE_PATH"
hook_class_name = "HOOK_CLASS_NAME"
operator_class = ClassBlock(
name="CLASS", extend_class="EXTEND_CLASS", methods_blocks=[]
)
operator_test_generator.create_operator_test_class_blocks(
operator_module_path=operator_module_path,
hook_class_name=hook_class_name,
operator_class=operator_class,
)
mock_create_assert_call_method_block.assert_called_once_with(
operator_module_path, hook_class_name, operator_class
)
def generate_class_block(client_info: ClientInfo,
integration: Integration) -> ClassBlock:
ctor_method = generate_ctor_method_block(client_info, integration)
get_conn_method = generate_get_conn_method_block(integration)
hook_methods = [ctor_method, get_conn_method]
for info in client_info.action_methods.values():
method_block = generate_method_block(
info, path_infos=client_info.path_methods, integration=integration)
hook_methods.append(method_block)
class_block = ClassBlock(
name=f"{integration.class_prefix}Hook",
extend_class=
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
methods_blocks=hook_methods,
)
return class_block
def test_should_skip_get_conn(self, mock_create_operator_class_block):
method_a = self._create_method("get_conn")
method_b = self._create_method("METHOD_A")
hook_class_block = ClassBlock(
name="CLASS_PREFIXHook",
extend_class=
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
methods_blocks=[method_a, method_b],
)
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLOENT_PATH",
)
result = operator_generator.create_operator_class_blocks(
hook_class_block, integration)
self.assertEqual(["CLASS_A"], result)
def test_generate_class_block(
self,
mock_generate_method_block,
mock_generate_get_conn_method_block,
mock_generate_ctor_method_block,
):
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLOENT_PATH",
)
ctor_method = self._create_action_info("CTOR_")
update_instance_method = self._create_action_info("UPDATE_INSTANCE_")
client_info = ClientInfo(
ctor_method=ctor_method,
path_methods={},
action_methods={"update_instance": update_instance_method},
)
class_block = hook_generator.generate_class_block(
client_info=client_info, integration=integration)
self.assertEqual(
ClassBlock(
name="CLASS_PREFIXHook",
extend_class=
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
methods_blocks=["METHOD_CTOR", "METHOD_GET_CONN", "METHOD_A"],
),
class_block,
)
mock_generate_method_block.assert_any_call(update_instance_method,
path_infos={},
integration=integration)
mock_generate_ctor_method_block.assert_called_once_with(
client_info, integration)
mock_generate_get_conn_method_block.assert_called_once_with(
integration)
mock_generate_ctor_method_block.assert_called_once_with(
client_info, integration)
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_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(name):
return ClassBlock(name=name, extend_class="EXTEND_CLASS", methods_blocks=[])
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],
)