potiuk commented on a change in pull request #5770: [AIRFLOW-5162] GCS Hook Upload Method Improvement URL: https://github.com/apache/airflow/pull/5770#discussion_r323329874
########## File path: tests/contrib/hooks/test_gcs_hook.py ########## @@ -602,21 +604,103 @@ def test_upload(self, mock_service): ) @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) - def test_upload_gzip(self, mock_service): + def test_upload_file_gzip(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' - upload_method = mock_service.return_value.get_bucket.return_value \ + upload_method = mock_service.return_value.get_bucket.return_value\ .blob.return_value.upload_from_filename upload_method.return_value = None response = self.gcs_hook.upload(test_bucket, # pylint:disable=assignment-from-no-return test_object, - self.testfile.name, + filename=self.testfile.name, gzip=True) self.assertFalse(os.path.exists(self.testfile.name + '.gz')) self.assertIsNone(response) + @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + def test_upload_string_strdata(self, mock_service): + test_bucket = 'test_bucket' + test_object = 'test_object' + + upload_method = mock_service.return_value.bucket.return_value\ + .blob.return_value.upload_from_string + upload_method.return_value = None + + response = self.gcs_hook.upload(test_bucket, # pylint:disable=assignment-from-no-return + test_object, + data=self.testdata_str) + + self.assertIsNone(response) + upload_method.assert_called_once_with( + self.testdata_str, + content_type='text/plain' + ) + + @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + def test_upload_string_bytedata(self, mock_service): + test_bucket = 'test_bucket' + test_object = 'test_object' + + upload_method = mock_service.return_value.bucket.return_value\ + .blob.return_value.upload_from_string + upload_method.return_value = None + + response = self.gcs_hook.upload(test_bucket, # pylint:disable=assignment-from-no-return + test_object, + data=self.testdata_bytes) + + self.assertIsNone(response) + upload_method.assert_called_once_with( + self.testdata_bytes, + content_type='text/plain' + ) + + @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + def test_upload_string_gzip(self, mock_service): + test_bucket = 'test_bucket' + test_object = 'test_object' + + upload_method = mock_service.return_value.get_bucket.return_value\ + .blob.return_value.upload_from_string + upload_method.return_value = None + + response = self.gcs_hook.upload(test_bucket, # pylint:disable=assignment-from-no-return + test_object, + data=self.testdata_str, + gzip=True) + + self.assertIsNone(response) + + response = self.gcs_hook.upload(test_bucket, # pylint:disable=assignment-from-no-return + test_object, + data=self.testdata_bytes, + gzip=True) + Review comment: I think we need to run asserr_called_once_with here. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services