diff --git a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py index bc039cf5575cb..66172b4c40ead 100644 --- a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py @@ -17,6 +17,7 @@ # under the License. # import unittest +from tempfile import NamedTemporaryFile from unittest import mock import pandas as pd @@ -33,29 +34,29 @@ def test_execute(self, mock_mysql_hook, mock_s3_hook, temp_mock): query = "query" s3_bucket = "bucket" s3_key = "key" - filename = "file" test_df = pd.DataFrame({'a': '1', 'b': '2'}, index=[0, 1]) get_pandas_df_mock = mock_mysql_hook.return_value.get_pandas_df get_pandas_df_mock.return_value = test_df - temp_mock.return_value.__enter__.return_value.name = filename - - op = MySQLToS3Operator(query=query, - s3_bucket=s3_bucket, - s3_key=s3_key, - mysql_conn_id="mysql_conn_id", - aws_conn_id="aws_conn_id", - task_id="task_id", - pd_csv_kwargs={'index': False, 'header': False}, - dag=None - ) - op.execute(None) - mock_mysql_hook.assert_called_once_with(mysql_conn_id="mysql_conn_id") - mock_s3_hook.assert_called_once_with(aws_conn_id="aws_conn_id", verify=None) - - get_pandas_df_mock.assert_called_once_with(query) - - temp_mock.assert_called_once_with(mode='r+', suffix=".csv") - mock_s3_hook.return_value.load_file.assert_called_once_with(filename=filename, - key=s3_key, - bucket_name=s3_bucket) + with NamedTemporaryFile() as f: + temp_mock.return_value.__enter__.return_value.name = f.name + + op = MySQLToS3Operator(query=query, + s3_bucket=s3_bucket, + s3_key=s3_key, + mysql_conn_id="mysql_conn_id", + aws_conn_id="aws_conn_id", + task_id="task_id", + pd_csv_kwargs={'index': False, 'header': False}, + dag=None + ) + op.execute(None) + mock_mysql_hook.assert_called_once_with(mysql_conn_id="mysql_conn_id") + mock_s3_hook.assert_called_once_with(aws_conn_id="aws_conn_id", verify=None) + + get_pandas_df_mock.assert_called_once_with(query) + + temp_mock.assert_called_once_with(mode='r+', suffix=".csv") + mock_s3_hook.return_value.load_file.assert_called_once_with(filename=f.name, + key=s3_key, + bucket_name=s3_bucket)