def test_get_docker_image_with_custom_image_uri(
self, mock_api_client, mock_logger):
self.setup()
# Verify mocking is correct and mock img tag.
assert mock_api_client is docker.APIClient
assert mock_logger is containerize.logger
docker_client = mock_api_client.return_value
lcb = containerize.LocalContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.worker_config,
destination_dir="/my_app/temp/",
docker_config=docker_config.DockerConfig(image="gcr.io/test-name"),
)
lcb._get_file_path_map = mock.Mock()
lcb._get_file_path_map.return_value = {}
lcb.tar_file_path = ""
with mock.patch("builtins.open", mock.mock_open(read_data="")):
img_tag = lcb.get_docker_image()
self.assertEqual(img_tag, "gcr.io/test-name")
# Verify docker APIClient is invoked as expected.
self.assertEqual(mock_api_client.call_count, 1)
_, kwargs = mock_api_client.call_args
self.assertDictEqual(kwargs, {"version": "auto"})
# Verify APIClient().build is invoked as expected.
self.assertEqual(docker_client.build.call_count, 1)
_, kwargs = docker_client.build.call_args
expected = {
"path": ".",
"custom_context": True,
"encoding": "utf-8",
"tag": img_tag,
}
self.assertTrue(set(expected.items()).issubset(set(kwargs.items())))
# Verify APIClient().push is invoked as expected.
self.assertEqual(docker_client.push.call_count, 1)
args, kwargs = docker_client.push.call_args
self.assertListEqual(list(args), [img_tag])
self.assertDictEqual(kwargs, {"decode": True, "stream": True})
# Verify logger info calls.
self.assertEqual(mock_logger.info.call_count, 2)
mock_logger.info.assert_has_calls(
[
mock.call("Building Docker image: %s", img_tag),
mock.call("Publishing Docker image: %s", img_tag),
]
)
self.cleanup(lcb.docker_file_path)
def test_create_docker_file_with_docker_parent_image(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
docker_config=docker_config.DockerConfig(
parent_image="tensorflow/tensorflow:latest"),
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:latest\n",
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.py"]',
]
self.assert_docker_file(expected_docker_file_lines,
lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
def test_get_docker_image_cloud_build_with_cache(
self, mock_gcs_client, mock_discovery_build
):
self.setup()
lcb = containerize.CloudContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.worker_config,
destination_dir="/my_app/temp/",
docker_config=docker_config.DockerConfig(
image="gcr.io/test-name:1",
parent_image="tensorflow/tensorflow:latest",
cache_from="gcr.io/test-name:0",
image_build_bucket="test_gcs_bucket"),
)
# Mock tar file generation
lcb._get_file_path_map = mock.Mock()
lcb._get_file_path_map.return_value = {}
# Mock cloud build return value
proj_ret_val = mock_discovery_build.return_value.projects.return_value
builds_ret_val = proj_ret_val.builds.return_value
create_ret_val = builds_ret_val.create.return_value
create_ret_val.execute.return_value = {
"metadata": {"build": {"id": "test_build_id"}}
}
get_ret_val = builds_ret_val.get.return_value
get_ret_val.execute.return_value = {"status": "SUCCESS"}
# Get docker image
img_tag = lcb.get_docker_image()
self.assertEqual(img_tag, "gcr.io/test-name:1")
# Verify gcs get_bucket is invoked
client_ret_val = mock_gcs_client.return_value
self.assertEqual(client_ret_val.get_bucket.call_count, 1)
args, _ = client_ret_val.get_bucket.call_args
self.assertListEqual(list(args), ["test_gcs_bucket"])
# Verify that a new blob is created in the bucket
get_bucket_ret_val = client_ret_val.get_bucket.return_value
self.assertEqual(get_bucket_ret_val.blob.call_count, 1)
args, _ = get_bucket_ret_val.blob.call_args
self.assertLen(args, 1)
storage_object_name = args[0]
self.assertTrue(storage_object_name, "tf_cloud_train_tar_")
# Verify that tarfile is added to the blob
blob_ret_val = get_bucket_ret_val.blob.return_value
self.assertEqual(blob_ret_val.upload_from_filename.call_count, 1)
# Verify discovery API is invoked as expected
self.assertEqual(mock_discovery_build.call_count, 1)
args, kwargs = mock_discovery_build.call_args
self.assertListEqual(list(args), ["cloudbuild", "v1"])
self.assertDictEqual(
kwargs,
{
"cache_discovery": False,
"requestBuilder": google_api_client.TFCloudHttpRequest,
},
)
# Verify cloud build create is invoked as expected
self.assertEqual(builds_ret_val.create.call_count, 1)
_, kwargs = builds_ret_val.create.call_args
request_dict = {}
request_dict["projectId"] = self.project_id
request_dict["images"] = [[img_tag]]
request_dict["steps"] = [{
"name": "gcr.io/cloud-builders/docker",
"entrypoint": "bash",
"args": [
"-c",
"docker pull gcr.io/test-name:0 || exit 0",
],
}, {
"name": "gcr.io/cloud-builders/docker",
"args": [
"build",
"-t",
"gcr.io/test-name:1",
"--cache-from",
"gcr.io/test-name:0",
"."],
}]
request_dict["source"] = {
"storageSource": {
"bucket": "test_gcs_bucket",
"object": storage_object_name,
}
}
self.assertDictEqual(
kwargs,
{
"projectId": self.project_id,
"body": request_dict,
},
)
# Verify cloud build get is invoked as expected
self.assertEqual(builds_ret_val.get.call_count, 1)
_, kwargs = builds_ret_val.get.call_args
self.assertDictEqual(
kwargs,
{
"projectId": self.project_id,
"id": "test_build_id",
},
)
self.cleanup(lcb.docker_file_path)
def test_valid(self):
docker_config.DockerConfig(
parent_image="tensorflow/tensorflow:latest-gpu",
image_build_bucket="test-bucket",
image="gcr.io/test-project/tf_cloud_train:1",
cache_from="gcr.io/test-project/tf_cloud_train:0")
