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 create_file_block(
operator_file_block: FileBlock, hook_file_block: FileBlock, integration: Integration
) -> FileBlock:
logging.info("Start creating operators test file block")
operator_module_name, _ = operator_file_block.file_name.rsplit(".", 2) # drop ".py"
operator_module_path = f"airflow.contrib.operator.{operator_module_name}"
hook_class_name = hook_file_block.class_blocks[0].name
test_class_blocks = [
create_operator_test_class_blocks(
operator_module_path, hook_class_name, operator_class_block
)
for operator_class_block in operator_file_block.class_blocks
]
constants = create_constants(operator_file_block.class_blocks)
file_block: FileBlock = FileBlock(
file_name=f"test_{integration.file_prefix}_operator.py",
class_blocks=test_class_blocks,
constants=constants,
)
imports_statement_gather.update_imports_statements(file_block)
for operator_class_block in operator_file_block.class_blocks:
operator_class_path = f"{operator_module_path}.{operator_class_block.name}"
file_block.import_statement.add(operator_class_path)
logging.info("Finished creating operators test file block")
return file_block
def create_file_block(hook_file_block: FileBlock, client_info: ClientInfo,
integration: Integration) -> FileBlock:
logging.info("Start creating hook test file block")
hook_module_name, _ = hook_file_block.file_name.rsplit(".",
2) # drop ".py"
hook_class = hook_file_block.class_blocks[0]
hook_class_path = f"airflow.contrib.hooks.{hook_module_name}.{hook_class.name}"
class_block_with_default_project_id = generate_class_block_with_default_project_id(
hook_class_path, hook_class, client_info, integration)
class_block_without_default_project_id = generate_class_block_without_default_project_id(
hook_class_path, hook_class, client_info, integration)
file_block: FileBlock = FileBlock(
file_name=f"test_{integration.file_prefix}_hook.py",
class_blocks=[
class_block_with_default_project_id,
class_block_without_default_project_id,
],
constants=generate_constants(hook_class.methods_blocks,
client_info.action_methods.values()),
)
imports_statement_gather.update_imports_statements(file_block)
file_block.import_statement.add(hook_class_path)
logging.info("Finished creating hook test file block")
return file_block
def create_file_block(client_info: ClientInfo,
integration: Integration) -> FileBlock:
logging.info("Start creating hook block")
class_block = generate_class_block(client_info, integration)
file_block: FileBlock = FileBlock(
file_name=f"{integration.file_prefix}_hook.py",
class_blocks=[class_block])
imports_statement_gather.update_imports_statements(file_block)
logging.info("Finished creating hook block")
return file_block
def test_render_file_block(self, mock_render_template, mock_render_class_block):
file_block = FileBlock(
file_name="FILE_NAME",
class_blocks=["CLASS_A"],
import_statement={"kitty", "mouse"},
constants=["CONSTANT_A"],
)
result = brushes.render_file_block(file_block)
mock_render_template.assert_called_once_with(
constants=["CONSTANT_A"],
class_blocks=["TEMPLATE_CLASS_BLOCK"],
import_statement={"kitty", "mouse"},
template_name="file_block.py.tpl",
)
self.assertEqual("TEMPLATE", result)
def create_file_block(integration: Integration,
hook_file_block: FileBlock) -> FileBlock:
logging.info("Start creating operators file block")
hook_module_name, _ = hook_file_block.file_name.rsplit(".",
2) # drop ".py"
hook_class = hook_file_block.class_blocks[0]
hook_class_path = f"airflow.contrib.hooks.{hook_module_name}.{hook_class.name}"
operator_class_blocks = create_operator_class_blocks(
hook_class, integration)
file_block: FileBlock = FileBlock(
file_name=f"{integration.file_prefix}_operator.py",
class_blocks=operator_class_blocks,
)
imports_statement_gather.update_imports_statements(file_block)
file_block.import_statement.add(hook_class_path)
logging.info("Finished creating operators file block")
return file_block
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_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
)