"""Unit tests for snowflake_snapshot_mechanical_transactions_task.""" from unittest.mock import patch from lib import config from tasks.accounting_period_mechanicals.snowflake_snapshot_mech_txns import \ snowflake_snapshot_mechanical_transactions_task import_path = 'tasks.accounting_period_mechanicals.snowflake_snapshot_mech_txns' @patch(f'{import_path}.RoyaltySnowflakeHook') @patch(f'{import_path}.get_accounting_period_details') @patch(f'{import_path}.shared_helpers') @patch(f'{import_path}.helpers') @patch(f'{import_path}.insert_snapshot_mechanical_transactions') def test_snowflake_snapshot_mechanical_transactions_task_success( mock_template, mock_helpers, mock_shared_helpers, mock_get_accounting_period_details, mock_hook, mock_accounting_period_mechanicals_event, mock_accounting_period_mechanicals_dag_run ): """Test task that snapshots sales that are eligible for mechanical deductions.""" accounting_period_id = mock_accounting_period_mechanicals_event['target_id'] statement_period_id = \ mock_accounting_period_mechanicals_event['statement_period_id'] accounting_period = { 'accounting_period_id': accounting_period_id, 'statement_period_id': statement_period_id } sales_file_ids = [4, 5, 6] mock_helpers.get_event_from_params.return_value.target_id = accounting_period_id mock_get_accounting_period_details.return_value = accounting_period mock_shared_helpers.get_sales_file_ids.return_value = sales_file_ids mock_template.return_value.render.return_value = \ 'INSERT INTO SNAPSHOT_MECHANICAL_TRANSACTIONS' mock_hook.return_value.run.return_value = True snowflake_snapshot_mechanical_transactions_task( mock_accounting_period_mechanicals_dag_run ) mock_helpers.get_event_from_params.assert_called_once_with( mock_accounting_period_mechanicals_dag_run ) mock_get_accounting_period_details.assert_called_once_with(accounting_period_id) mock_shared_helpers.get_sales_file_ids.assert_called_once_with(accounting_period_id) mock_hook.assert_called_once_with(snowflake_conn_id=config.SNOWFLAKE_CONN_NAME) mock_template.return_value.render.assert_called_once_with( accounting_period_id=accounting_period_id, sales_file_ids=sales_file_ids, schema=config.OWS_ENV, statement_period_id=statement_period_id )