"""Unit tets for snowflake_flatten_contract_json NR task.""" from unittest.mock import patch from lib import config from tasks.accounting_run_calculate_nr.snowflake_flatten_contract_json \ import insert_flat_contract_task import_path = 'tasks.accounting_run_calculate_nr.snowflake_flatten_contract_json' @patch(f'{import_path}.RoyaltySnowflakeHook') @patch(f'{import_path}.helpers') @patch(f'{import_path}.insert_flat_nr_contract_data_template') def test_snowflake_insert_flat_nr_contract_task( mock_template, mock_helpers, mock_hook, mock_accounting_run_calculate_nr_dag_run ): """Test that task renders an insert query.""" accounting_run_id = 123 mock_helpers.get_event_from_params.return_value.target_id = accounting_run_id mock_template.return_value.render.return_value = \ 'INSERT FLAT CONTRACT, CONTRACT_TERM DATA INTO CONTRACT_DENORMALIZED_NR' mock_hook.return_value.run.return_value = True insert_flat_contract_task(mock_accounting_run_calculate_nr_dag_run) mock_helpers.get_event_from_params.assert_called_once_with( mock_accounting_run_calculate_nr_dag_run ) mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME) mock_template.return_value.render.assert_called_once_with( accounting_run_id=accounting_run_id, schema=config.OWS_ENV ) mock_hook.return_value.run.assert_called_once() @patch(f'{import_path}.RoyaltySnowflakeHook') @patch(f'{import_path}.helpers') @patch(f'{import_path}.insert_flat_nr_contract_contributor_only_data_template') def test_snowflake_insert_flat_nr_contract_contributor_only_task( mock_template, mock_helpers, mock_hook, mock_accounting_run_calculate_nr_dag_run ): """Test that task renders an insert query.""" accounting_run_id = 123 mock_helpers.get_event_from_params.return_value.target_id = accounting_run_id mock_template.return_value.render.return_value = \ 'INSERT CONTRIBUTOR CONTRACT DATA INTO CONTRACT_DENORMALIZED_NR_CONTRIB_ONLY' mock_hook.return_value.run.return_value = True insert_flat_contract_task( mock_accounting_run_calculate_nr_dag_run, is_contributor_only=True ) mock_helpers.get_event_from_params.assert_called_once_with( mock_accounting_run_calculate_nr_dag_run ) mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME) mock_template.return_value.render.assert_called_once_with( accounting_run_id=accounting_run_id, schema=config.OWS_ENV ) mock_hook.return_value.run.assert_called_once()