diff --git a/test/sample-test/sample_test_launcher.py b/test/sample-test/sample_test_launcher.py index d7c9a31b6a1..7098e36c401 100644 --- a/test/sample-test/sample_test_launcher.py +++ b/test/sample-test/sample_test_launcher.py @@ -27,6 +27,7 @@ from constants import PAPERMILL_ERR_MSG, BASE_DIR, TEST_DIR, SCHEMA_CONFIG, CONFIG_DIR, DEFAULT_CONFIG from check_notebook_results import NoteBookChecker +from kfp.compiler._gcs_helper import GCSHelper from run_sample_test import PySampleChecker @@ -56,14 +57,9 @@ def _copy_result(self): """ Copy generated sample test result to gcs, so that Prow can pick it. """ print('Copy the test results to GCS %s/' % self._results_gcs_dir) - # TODO(Issue#2076): Currently switch to gsutil. Switch back to upload_blob - # when it is fixed. - subprocess.call([ - 'gsutil', - 'cp', + GCSHelper.upload_gcs_file( self._sample_test_result, - os.path.join(self._results_gcs_dir, self._sample_test_result) - ]) + os.path.join(self._results_gcs_dir, self._sample_test_result)) def _compile(self): diff --git a/test/sample-test/utils.py b/test/sample-test/utils.py index 27a87f88449..cfb090fbe9f 100644 --- a/test/sample-test/utils.py +++ b/test/sample-test/utils.py @@ -88,20 +88,3 @@ def file_injection(file_in, tmp_file_out, subs): os.rename(tmp_file_out, file_in) -# TODO(Issue#2076): Currently broken, further investigation needed. -def upload_blob(bucket_name, source_file_name, destination_blob_name): - """Uploads a file to the bucket.""" - storage_client = storage.Client() - try: - bucket = storage_client.get_bucket(bucket_name) - blob = bucket.blob(destination_blob_name) - blob.upload_from_filename(source_file_name) - except google.cloud.exceptions.GoogleCloudError as google_cloud_error: - raise RuntimeError( - 'Failure when uploading {}\n'.format(str(google_cloud_error))) - except google.cloud.exceptions.NotFound as not_found: - raise RuntimeError("Bucket not found: {}\n".format(str(not_found))) - else: - print('File {} uploaded to {}.'.format( - source_file_name, - destination_blob_name))