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 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_update_imports_statements(self, mock_gather_imports):
file_block = mock.MagicMock()
imports_statement_gather.update_imports_statements(file_block)
mock_gather_imports.assert_called_once_with(file_block)
file_block.import_statement.update.assert_called_once_with(["IMPORT_A"])