def test_create_streaming_dataflow_job_when_job_does_not_already_exist(
self):
"""Test that attempting to deploy an update to a Dataflow job when a job with the name of service does not
already exist results in the job deployment being retried with `update` set to `False`.
"""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION, temporary_directory)
deployer = DataflowDeployer(octue_configuration_path)
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.Topic",
return_value=Mock(
path="projects/my-project/topics/my-topic"),
):
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.DataflowRunner.run_pipeline",
side_effect=[ValueError, None],
) as mock_runner:
deployer.create_streaming_dataflow_job(
image_uri="my-image-uri", update=True)
# Check that the first attempt was to update the service.
first_attempt_options = mock_runner.mock_calls[0].kwargs[
"options"].get_all_options()
self.assertTrue(first_attempt_options["update"])
# Check that the second attempt was to create the service.
second_attempt_options = mock_runner.mock_calls[1].kwargs[
"options"].get_all_options()
self.assertFalse(second_attempt_options["update"])
def test_dataflow_prime_enabled_if_machine_type_not_specified(self):
"""Test that Dataflow Prime is enabled if the machine type is not specified in the octue configuration."""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_with_no_machine_type = copy.deepcopy(
OCTUE_CONFIGURATION)
del octue_configuration_with_no_machine_type["services"][0][
"machine_type"]
octue_configuration_path = self._create_octue_configuration_file(
octue_configuration_with_no_machine_type,
temporary_directory,
)
deployer = DataflowDeployer(octue_configuration_path)
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.Topic",
return_value=Mock(
path="projects/my-project/topics/my-topic"),
):
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.DataflowRunner.run_pipeline"
) as mock_run:
deployer.create_streaming_dataflow_job(
image_uri="my-image-uri")
self.assertEqual(
mock_run.call_args.kwargs["options"].get_all_options()
["dataflow_service_options"],
["enable_prime"],
)
def test_updating_streaming_dataflow_job(self):
"""Test updating an existing streaming dataflow job."""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION, temporary_directory)
deployer = DataflowDeployer(octue_configuration_path)
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.Topic",
return_value=Mock(
path="projects/my-project/topics/my-topic"),
):
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.DataflowRunner.run_pipeline"
) as mock_run:
deployer.create_streaming_dataflow_job(
image_uri="my-image-uri", update=True)
options = mock_run.call_args.kwargs["options"].get_all_options()
self.assertTrue(options["update"])
self.assertTrue(options["streaming"])
self.assertIsNone(options["dataflow_service_options"])
self.assertEqual(options["project"], SERVICE["project_name"])
self.assertEqual(options["job_name"], SERVICE["name"])
self.assertEqual(options["temp_location"],
DEFAULT_DATAFLOW_TEMPORARY_FILES_LOCATION)
self.assertEqual(options["region"], SERVICE["region"])
self.assertEqual(options["sdk_container_image"], "my-image-uri")
self.assertEqual(options["setup_file"], DEFAULT_SETUP_FILE_PATH)
def test_generate_cloud_build_configuration(self):
"""Test that a correct Google Cloud Build configuration is generated from the given `octue.yaml` file."""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION, temporary_directory)
deployer = DataflowDeployer(octue_configuration_path)
deployer._generate_cloud_build_configuration()
self.assertEqual(deployer.generated_cloud_build_configuration,
EXPECTED_CLOUD_BUILD_CONFIGURATION)
def test_deployment_error_raised_if_dataflow_job_already_exists(self):
"""Test that a deployment error is raised if a Dataflow job already exists with the same name as the service."""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION, temporary_directory)
deployer = DataflowDeployer(octue_configuration_path)
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.Topic",
return_value=Mock(
path="projects/my-project/topics/my-topic"),
):
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.DataflowRunner.run_pipeline",
side_effect=DataflowJobAlreadyExistsError(),
):
with self.assertRaises(DeploymentError):
deployer.create_streaming_dataflow_job(
image_uri="my-image-uri", update=True)
def dataflow(service_config, no_cache, update, dataflow_job_only, image_uri):
"""Deploy a python app to Google Dataflow as an Octue service or digital twin."""
if bool(importlib.util.find_spec("apache_beam")):
# Import the Dataflow deployer only if the `apache-beam` package is available (due to installing `octue` with
# the `dataflow` extras option).
from octue.cloud.deployment.google.dataflow.deployer import DataflowDeployer
else:
raise ImportWarning(
"To use this CLI command, you must install `octue` with the `dataflow` option e.g. "
"`pip install octue[dataflow]`")
deployer = DataflowDeployer(service_config)
if dataflow_job_only:
deployer.create_streaming_dataflow_job(image_uri=image_uri,
update=update)
return
deployer.deploy(no_cache=no_cache, update=update)
def test_deploy_with_cloud_build_file_provided(self):
"""Test deploying to Dataflow with a `cloudbuild.yaml` path provided in the `octue.yaml` file"""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION_WITH_CLOUD_BUILD_PATH,
temporary_directory,
)
deployer = DataflowDeployer(octue_configuration_path,
image_uri_template="blah")
with patch("subprocess.run",
return_value=Mock(returncode=0)) as mock_run:
mock_build_id = "my-build-id"
with patch(
"json.loads",
return_value={
"metadata": {
"build": {
"images": [deployer.image_uri_template],
"id": mock_build_id
}
},
"status": "SUCCESS",
},
):
with patch(
"octue.cloud.deployment.google.dataflow.pipeline.DataflowRunner"
):
deployer.deploy()
# Test the build trigger creation request.
self.assertEqual(
mock_run.call_args_list[0].args[0],
EXPECTED_BUILD_TRIGGER_CREATION_COMMAND + [
f"--build-config={OCTUE_CONFIGURATION_WITH_CLOUD_BUILD_PATH['services'][0]['cloud_build_configuration_path']}"
],
)
# Test the build trigger run request.
self.assertEqual(
mock_run.call_args_list[1].args[0],
[
"gcloud",
f"--project={OCTUE_CONFIGURATION_WITH_CLOUD_BUILD_PATH['services'][0]['project_name']}",
"--format=json",
"beta",
"builds",
"triggers",
"run",
OCTUE_CONFIGURATION_WITH_CLOUD_BUILD_PATH["services"][0]
["name"],
"--branch=my-branch",
],
)
# Test waiting for the build trigger run to complete.
self.assertEqual(
mock_run.call_args_list[2].args[0],
[
"gcloud",
f'--project={SERVICE["project_name"]}',
"--format=json",
"builds",
"describe",
mock_build_id,
],
)
def test_deploy(self):
"""Test that the build trigger creation and run are requested correctly."""
with tempfile.TemporaryDirectory() as temporary_directory:
octue_configuration_path = self._create_octue_configuration_file(
OCTUE_CONFIGURATION, temporary_directory)
deployer = DataflowDeployer(octue_configuration_path)
with patch("subprocess.run",
return_value=Mock(returncode=0)) as mock_run:
mock_build_id = "my-build-id"
with patch(
"json.loads",
return_value={
"metadata": {
"build": {
"images": [deployer.image_uri_template],
"id": mock_build_id
}
},
"status": "SUCCESS",
},
):
temporary_file = tempfile.NamedTemporaryFile(delete=False)
with patch("tempfile.NamedTemporaryFile",
return_value=temporary_file):
deployer.deploy()
# Test the build trigger creation request.
self.assertEqual(
mock_run.call_args_list[0].args[0],
EXPECTED_BUILD_TRIGGER_CREATION_COMMAND +
[f"--inline-config={temporary_file.name}"],
)
# Test the build trigger run request.
self.assertEqual(
mock_run.call_args_list[1].args[0],
[
"gcloud",
f"--project={SERVICE['project_name']}",
"--format=json",
"beta",
"builds",
"triggers",
"run",
SERVICE["name"],
"--branch=my-branch",
],
)
# Test waiting for the build trigger run to complete.
self.assertEqual(
mock_run.call_args_list[2].args[0],
[
"gcloud",
f'--project={SERVICE["project_name"]}',
"--format=json",
"builds",
"describe",
mock_build_id,
],
)