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 test_create_file_block(
self,
mock_imports_statement_gather,
mock_create_operator_test_class_blocks,
mock_create_constants,
):
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLOENT_PATH",
)
operator_a = self._create_operator_class_block("OperatorA")
operator_b = self._create_operator_class_block("OperatorB")
operator_file_block = FileBlock(
file_name="FILE_PREFFIX_operator.py",
class_blocks=[operator_a, operator_b],
import_statement={"IMPORT_A"},
)
hook_file_block = FileBlock(
file_name="FILE_PREFFIX_hook.py",
class_blocks=[operator_a, operator_b],
import_statement={"IMPORT_A"},
)
result = operator_test_generator.create_file_block(
operator_file_block, hook_file_block, integration
)
self.assertEqual(
FileBlock(
file_name="test_FILE_PREFFIX_operator.py",
class_blocks=["CLASS_A", "CLASS_B"],
import_statement={
"typing.Tuple",
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
"typing.Sequence",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_default_project_id",
"tests.contrib.utils.base_gcp_mock.GCP_PROJECT_ID_HOOK_UNIT_TEST",
"airflow.utils.decorators.apply_defaults",
"airflow.AirflowException",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_no_default_project_id",
"typing.Optional",
"typing.Dict",
"unittest.TestCase",
"airflow.contrib.operator.FILE_PREFFIX_operator.OperatorB",
"google.api_core.retry.Retry",
"unittest.mock",
"google.cloud.redis_v1beta1.CloudRedisClient",
"airflow.contrib.operator.FILE_PREFFIX_operator.OperatorA",
"typing.Union",
},
constants="CONSTANTS",
),
result,
)
mock_imports_statement_gather.update_imports_statements.assert_called_once_with(
result
)
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_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 main() -> None:
logging.basicConfig(level=logging.INFO)
integration_info: Integration = Integration(
service_name="Cloud Memorystore",
class_prefix="CloudMemorystore",
file_prefix="gcp_cloud_memorystore",
client_path="google.cloud.redis_v1.CloudRedisClient",
)
client_info: ClientInfo = client_parser.parse_cloud_client(
integration_info)
file_blocks: List[FileBlock] = block_generator.generate_file_blocks(
client_info=client_info, integration=integration_info)
block_renderer.create_files_from_file_blocks(file_blocks)
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_create_file_block(self, mock_generate_class_block,
mock_imports_statement_gather):
client_info = "CLIENT_INFO"
integration = Integration(
service_name="SERVICE_NAME",
class_prefix="CLASS_PREFIX",
file_prefix="FILE_PREFFIX",
client_path="CLOENT_PATH",
)
file_block = hook_generator.create_file_block(client_info, integration)
mock_generate_class_block.assert_called_once_with(
client_info, integration)
mock_imports_statement_gather.update_imports_statements.assert_called_once_with(
mock.ANY)
self.assertEqual(
FileBlock(
file_name="FILE_PREFFIX_hook.py",
class_blocks=["CLASS_A"],
import_statement={
"typing.Union",
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
"airflow.utils.decorators.apply_defaults",
"typing.Dict",
"typing.Tuple",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_no_default_project_id",
"google.cloud.redis_v1beta1.CloudRedisClient",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_default_project_id",
"airflow.AirflowException",
"google.api_core.retry.Retry",
"unittest.mock",
"typing.Optional",
"unittest.TestCase",
"tests.contrib.utils.base_gcp_mock.GCP_PROJECT_ID_HOOK_UNIT_TEST",
"typing.Sequence",
},
constants=[],
),
file_block,
)
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_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_file_block(
self,
mock_generate_class_block_without_default_project_id,
mock_imports_statement_gather,
mock_generate_class_block_with_default_project_id,
mock_generate_constants,
):
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"},
)
client_info = ClientInfo(
ctor_method=None,
path_methods={},
action_methods={"METHOD_A": mock.MagicMock(), "METHOD_B": mock.MagicMock()},
)
result = hook_test_generator.create_file_block(
hook_file_block=hook_file_block,
integration=integration,
client_info=client_info,
)
self.assertEqual(
FileBlock(
file_name="test_FILE_PREFFIX_hook.py",
class_blocks=["CLASS_A", "CLASS_B"],
import_statement={
"typing.Union",
"airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook",
"airflow.utils.decorators.apply_defaults",
"airflow.contrib.hooks.FILE_PREFFIX_hook.CLASS_PREFIXHook",
"typing.Dict",
"typing.Tuple",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_no_default_project_id",
"google.cloud.redis_v1beta1.CloudRedisClient",
"tests.contrib.utils.base_gcp_mock.mock_base_gcp_hook_default_project_id",
"airflow.AirflowException",
"google.api_core.retry.Retry",
"unittest.mock",
"typing.Optional",
"unittest.TestCase",
"tests.contrib.utils.base_gcp_mock.GCP_PROJECT_ID_HOOK_UNIT_TEST",
"typing.Sequence",
},
constants="CONSTANTS",
),
result,
)
mock_generate_class_block_without_default_project_id.assert_called_once_with(
"airflow.contrib.hooks.FILE_PREFFIX_hook.CLASS_PREFIXHook",
hook_class_block,
client_info,
integration,
)
mock_generate_class_block_with_default_project_id.assert_called_once_with(
"airflow.contrib.hooks.FILE_PREFFIX_hook.CLASS_PREFIXHook",
hook_class_block,
client_info,
integration,
)
mock_imports_statement_gather.update_imports_statements.assert_called_once_with(
result
)
mock_generate_constants.assert_called_once_with(
["METHOD_CTOR", "METHOD_GET_CONN", "METHOD_A"], mock.ANY
)