Skip to content

Commit

Permalink
Fixes location of temporary file created in tests (#9403)
Browse files Browse the repository at this point in the history
Solves small Issue created in a60f589

The file was created in the sources 'file' when run from IDE which
might be accidentally committed. This change prevents it.
  • Loading branch information
potiuk authored Jun 19, 2020
1 parent d7ef352 commit 760bee8
Showing 1 changed file with 23 additions and 22 deletions.
45 changes: 23 additions & 22 deletions tests/providers/amazon/aws/transfers/test_mysql_to_s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
# under the License.
#
import unittest
from tempfile import NamedTemporaryFile
from unittest import mock

import pandas as pd
Expand All @@ -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)

0 comments on commit 760bee8

Please sign in to comment.