"""Test snowflake_copy_contract_transactions.""" from unittest.mock import MagicMock, patch from lib import config from tasks.accounting_run_commit.snowflake_copy_contract_transactions import\ copy_contract_transactions @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.ows.get_accounting_period_details') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.ows.get_accounting_run_details') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.RoyaltySnowflakeHook') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.copy_contract_transaction_distro_template') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.get_event_from_params') # noqa: E501 def test_copy_contract_transaction_distribution( mock_get_event: MagicMock, mock_template: MagicMock, mock_hook: MagicMock, mock_ows_get_accounting_run_details: MagicMock, mock_ows_get_accounting_period_details: MagicMock, mock_commit_event, mock_commit_dag_run, mock_accounting_run, ): """Test using snowflake hook to copy distribution contract transactions from staging to 'final' table.""" # noqa: E501 accounting_run_id = mock_commit_event.get('target_id') mock_accounting_period = { 'accounting_period_id': 1, 'accounting_period_status': 'open', 'contract_type': 'distribution' } mock_ows_get_accounting_run_details.return_value = mock_accounting_run mock_ows_get_accounting_period_details.return_value = mock_accounting_period mock_get_event.return_value.target_id = accounting_run_id mock_template.return_value.render.return_value = 'COPY STATEMENT' mock_hook.return_value.run.return_value = True copy_contract_transactions(mock_commit_dag_run) mock_get_event.assert_called_once_with(mock_commit_dag_run) mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME) mock_template.assert_called_once() mock_template.return_value.render.assert_called_once_with( accounting_run_id=accounting_run_id, environment=config.OWS_ENV, ) mock_hook.return_value.run.assert_called_once() mock_ows_get_accounting_run_details.assert_called_once_with(accounting_run_id) mock_ows_get_accounting_period_details.assert_called_once_with( mock_accounting_period['accounting_period_id']) @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.ows.get_accounting_period_details') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.ows.get_accounting_run_details') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.RoyaltySnowflakeHook') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.copy_contract_transaction_nr_template') # noqa: E501 @patch('tasks.accounting_run_commit.snowflake_copy_contract_transactions.get_event_from_params') # noqa: E501 def test_copy_contract_transaction_nr( mock_get_event: MagicMock, mock_template: MagicMock, mock_hook: MagicMock, mock_ows_get_accounting_run_details: MagicMock, mock_ows_get_accounting_period_details: MagicMock, mock_commit_event, mock_commit_dag_run, mock_accounting_run, ): """Test using snowflake hook to copy NR contract transactions from staging to 'final' table.""" # noqa: E501 accounting_run_id = mock_commit_event.get('target_id') mock_accounting_period = { 'accounting_period_id': 1, 'accounting_period_status': 'open', 'contract_type': 'neighbouring_rights' } mock_ows_get_accounting_run_details.return_value = mock_accounting_run mock_ows_get_accounting_period_details.return_value = mock_accounting_period mock_get_event.return_value.target_id = accounting_run_id mock_template.return_value.render.return_value = 'COPY STATEMENT' mock_hook.return_value.run.return_value = True copy_contract_transactions(mock_commit_dag_run) mock_get_event.assert_called_once_with(mock_commit_dag_run) mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME) mock_template.assert_called_once() mock_template.return_value.render.assert_called_once_with( accounting_run_id=accounting_run_id, environment=config.OWS_ENV, ) mock_hook.return_value.run.assert_called_once() mock_ows_get_accounting_run_details.assert_called_once_with(accounting_run_id) mock_ows_get_accounting_period_details.assert_called_once_with( mock_accounting_period['accounting_period_id'])