diff --git a/airflow/providers/amazon/aws/operators/s3_file_transform.py b/airflow/providers/amazon/aws/operators/s3_file_transform.py index 39788918a8bdc..4f0e626fc37e8 100644 --- a/airflow/providers/amazon/aws/operators/s3_file_transform.py +++ b/airflow/providers/amazon/aws/operators/s3_file_transform.py @@ -167,7 +167,7 @@ def execute(self, context): self.log.info("Uploading transformed file to S3") f_dest.flush() dest_s3.load_file( - filename=f_dest.name, + filename=f_dest.name if self.transform_script else f_source.name, key=self.dest_s3_key, replace=self.replace ) diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py b/tests/providers/amazon/aws/operators/test_s3_file_transform.py index 602eef1ab3e57..a1bf7beeb5ed0 100644 --- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py +++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py @@ -36,6 +36,11 @@ class TestS3FileTransformOperator(unittest.TestCase): def setUp(self): + self.content = b"input" + self.bucket = "bucket" + self.input_key = "foo" + self.output_key = "bar" + self.bio = io.BytesIO(self.content) self.tmp_dir = mkdtemp(prefix='test_tmpS3FileTransform_') self.transform_script = os.path.join(self.tmp_dir, "transform.py") os.mknod(self.transform_script) @@ -123,6 +128,10 @@ def test_execute_with_select_expression(self, mock_select_key): expression=select_expression ) + conn = boto3.client('s3') + result = conn.get_object(Bucket=self.bucket, Key=self.output_key) + self.assertEqual(self.content, result['Body'].read()) + @staticmethod def mock_process(mock_popen, return_code=0, process_output=None): process = mock_popen.return_value @@ -130,19 +139,13 @@ def mock_process(mock_popen, return_code=0, process_output=None): process.wait.return_value = None process.returncode = return_code - @staticmethod - def s3_paths(): - bucket = "bucket" - input_key = "foo" - output_key = "bar" - bio = io.BytesIO(b"input") - + def s3_paths(self): conn = boto3.client('s3') - conn.create_bucket(Bucket=bucket) - conn.upload_fileobj(Bucket=bucket, Key=input_key, Fileobj=bio) + conn.create_bucket(Bucket=self.bucket) + conn.upload_fileobj(Bucket=self.bucket, Key=self.input_key, Fileobj=self.bio) s3_url = "s3://{0}/{1}" - input_path = s3_url.format(bucket, input_key) - output_path = s3_url.format(bucket, output_key) + input_path = s3_url.format(self.bucket, self.input_key) + output_path = s3_url.format(self.bucket, self.output_key) return input_path, output_path