Skip to content

Commit

Permalink
fix unit tests and address some comments (kubeflow#1892)
Browse files Browse the repository at this point in the history
  • Loading branch information
gaoning777 authored and k8s-ci-robot committed Aug 20, 2019
1 parent 6a7b28f commit 79c7bda
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
4 changes: 2 additions & 2 deletions sdk/python/kfp/compiler/_container_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def __init__(self, gcs_staging=None, gcr_image_tag=None, namespace=None):
try:
gcs_bucket = self._get_project_id()
except:
raise ValueError('Please provide the gcr_staging.')
raise ValueError('Cannot get the Google Cloud project ID, please specify the gcs_staging argument.')
self._gcs_staging = 'gs://' + gcs_bucket + '/' + GCS_STAGING_BLOB_DEFAULT_PREFIX
else:
from pathlib import PurePath
Expand Down Expand Up @@ -160,7 +160,7 @@ def build(self, local_dir, docker_filename, target_image=None, timeout=1000):
target_image (str): the target image tag to push the final image.
timeout (int): time out in seconds. Default: 1000
"""
target_image = self._gcr_image_tag if target_image is None else target_image
target_image = target_image or self._gcr_image_tag
# Prepare build context
with tempfile.TemporaryDirectory() as local_build_dir:
from ._gcs_helper import GCSHelper
Expand Down
5 changes: 3 additions & 2 deletions sdk/python/tests/compiler/container_builder_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from kfp.compiler._component_builder import ContainerBuilder

GCS_BASE = 'gs://kfp-testing/'
GCR_IMAGE_TAG = 'gcr.io/kfp-testing/image'

@mock.patch('kfp.compiler._gcs_helper.GCSHelper')
class TestContainerBuild(unittest.TestCase):
Expand All @@ -39,7 +40,7 @@ def test_wrap_dir_in_tarball(self, mock_gcshelper):
f.write('temporary file two content')

# check
builder = ContainerBuilder(gcs_staging=GCS_BASE, namespace='')
builder = ContainerBuilder(gcs_staging=GCS_BASE, gcr_image_tag=GCR_IMAGE_TAG, namespace='')
builder._wrap_dir_in_tarball(temp_tarball, test_data_dir)
self.assertTrue(os.path.exists(temp_tarball))
with tarfile.open(temp_tarball) as temp_tarball_handle:
Expand All @@ -57,7 +58,7 @@ def test_generate_kaniko_yaml(self, mock_gcshelper):
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')

# check
builder = ContainerBuilder(gcs_staging=GCS_BASE, namespace='default')
builder = ContainerBuilder(gcs_staging=GCS_BASE, gcr_image_tag=GCR_IMAGE_TAG, namespace='default')
generated_yaml = builder._generate_kaniko_spec(docker_filename='dockerfile',
context='gs://mlpipeline/kaniko_build.tar.gz', target_image='gcr.io/mlpipeline/kaniko_image:latest')
with open(os.path.join(test_data_dir, 'kaniko.basic.yaml'), 'r') as f:
Expand Down

0 comments on commit 79c7bda

Please sign in to comment.