Use GCSHelper to upload test result (#2078)

* Add GCSHelper in sample_test_launcher

* Remove redundant stuff

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):
 

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))