diff --git a/workflows/cloud-client/cancel_execution.py b/workflows/cloud-client/cancel_execution.py new file mode 100644 index 00000000000..598ce67a85a --- /dev/null +++ b/workflows/cloud-client/cancel_execution.py @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_cancel_execution] +import sys + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 + + +def cancel_execution( + project_id: str, location_id: str, workflow_id: str, execution_id: str +) -> None: + """Cancels an execution that's still running. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'myFirstWorkflow'). + execution_id: The ID of the execution. + """ + client = executions_v1.ExecutionsClient() + name = ( + f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + f"/executions/{execution_id}" + ) + + try: + execution = client.cancel_execution(name=name) + except exceptions.NotFound as err: + print( + f"error: execution '{execution_id}' not found " + f"for workflow '{workflow_id}' in location '{location_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to cancel execution '{execution_id}' " + f"for workflow '{workflow_id}' in location '{location_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {execution.name}") + print(f"End Time: {execution.end_time}") + print(f"State: {execution.state.name}") + + +# [END workflows_cancel_execution] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=cancel_execution.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + parser.add_argument( + "--execution-id", + default=os.environ.get("EXECUTION_ID", "example-execution-id"), + help="ID of the execution (default: $EXECUTION_ID)", + ) + args = parser.parse_args() + + cancel_execution( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + execution_id=args.execution_id, + ) diff --git a/workflows/cloud-client/conftest.py b/workflows/cloud-client/conftest.py deleted file mode 100644 index a5d52654479..00000000000 --- a/workflows/cloud-client/conftest.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -import time -import uuid - -from google.cloud import workflows_v1 - -import pytest - - -PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -LOCATION = "us-central1" -WORKFLOW_ID_BASE = "myFirstWorkflow" - - -def workflow_exists(client: workflows_v1.WorkflowsClient, workflow_id: str) -> bool: - """Returns True if the workflow exists in this project.""" - try: - workflow_name = client.workflow_path( - PROJECT_ID, LOCATION, workflow_id - ) - client.get_workflow(request={"name": workflow_name}) - return True - except Exception as e: - print(f"Workflow doesn't exist: {e}") - return False - - -@pytest.fixture(scope="module") -def client() -> str: - assert PROJECT_ID, "'GOOGLE_CLOUD_PROJECT' environment variable not set." - workflows_client = workflows_v1.WorkflowsClient() - return workflows_client - - -@pytest.fixture(scope="module") -def project_id() -> str: - return PROJECT_ID - - -@pytest.fixture(scope="module") -def location() -> str: - return LOCATION - - -@pytest.fixture(scope="function") -def workflow_id(client: workflows_v1.WorkflowsClient) -> str: - workflow_id_str = f"{WORKFLOW_ID_BASE}_{uuid.uuid4()}" - - creating_workflow = False - backoff_delay = 1 # Start wait with delay of 1 second. - - # Create the workflow if it doesn't exist. - while not workflow_exists(client, workflow_id_str): - if not creating_workflow: - # Create the workflow. - workflow_file = open("myFirstWorkflow.workflows.yaml").read() - - parent = client.common_location_path(PROJECT_ID, LOCATION) - - client.create_workflow( - request={ - "parent": parent, - "workflow_id": workflow_id_str, - "workflow": { - "name": workflow_id_str, - "source_contents": workflow_file - }, - } - ) - - creating_workflow = True - - # Wait until the workflow is created. - print("- Waiting for the Workflow to be created...") - time.sleep(backoff_delay) - # Double the delay to provide exponential backoff. - backoff_delay *= 2 - - yield workflow_id_str - - # Delete the workflow. - workflow_full_name = client.workflow_path( - PROJECT_ID, LOCATION, workflow_id_str - ) - - client.delete_workflow( - request={ - "name": workflow_full_name, - } - ) diff --git a/workflows/cloud-client/create_workflow.py b/workflows/cloud-client/create_workflow.py new file mode 100644 index 00000000000..86a53cc618d --- /dev/null +++ b/workflows/cloud-client/create_workflow.py @@ -0,0 +1,120 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_create_workflow] +import sys + +from google.api_core import exceptions +from google.cloud import workflows_v1 + + +def create_workflow( + project_id: str, location_id: str, workflow_id: str, service_account_email: str +) -> None: + """Creates a workflow with logging, labels, and environment variables. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID for the new workflow (for example, 'example-workflow'). + service_account_email: The service account that the workflow runs as. + """ + client = workflows_v1.WorkflowsClient() + parent = f"projects/{project_id}/locations/{location_id}" + + # The workflow definition, in YAML. It returns a greeting from an + # environment variable that's set on the workflow below. + source_contents = """ +main: + steps: + - greet: + return: ${sys.get_env("EXAMPLE_GREETING")} +""" + workflow = workflows_v1.Workflow( + source_contents=source_contents, + service_account=service_account_email, + # Logs each call that fails, without logging successful calls. + call_log_level=workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY, + labels={"environment": "example"}, + user_env_vars={"EXAMPLE_GREETING": "example-greeting"}, + ) + + try: + operation = client.create_workflow( + parent=parent, workflow=workflow, workflow_id=workflow_id + ) + print("Waiting for the workflow to deploy...") + response = operation.result(timeout=300) + except exceptions.AlreadyExists as err: + print( + f"error: workflow '{workflow_id}' already exists " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to create workflow '{workflow_id}' " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Call Log Level: {response.call_log_level.name}") + print(f"Revision ID: {response.revision_id}") + print(f"State: {response.state.name}") + + +# [END workflows_create_workflow] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=create_workflow.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "example-workflow"), + help="ID for the new workflow (default: $WORKFLOW_ID)", + ) + parser.add_argument( + "--service-account-email", + default=os.environ.get( + "SERVICE_ACCOUNT_EMAIL", + "example-runner@example-project-id.iam.gserviceaccount.com", + ), + help="Service account the workflow runs as (default: $SERVICE_ACCOUNT_EMAIL)", + ) + args = parser.parse_args() + + create_workflow( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + service_account_email=args.service_account_email, + ) diff --git a/workflows/cloud-client/delete_workflow.py b/workflows/cloud-client/delete_workflow.py new file mode 100644 index 00000000000..bc5be0c2647 --- /dev/null +++ b/workflows/cloud-client/delete_workflow.py @@ -0,0 +1,86 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_delete_workflow] +import sys + +from google.api_core import exceptions +from google.cloud import workflows_v1 + + +def delete_workflow(project_id: str, location_id: str, workflow_id: str) -> None: + """Deletes a workflow and its executions. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'example-workflow'). + """ + client = workflows_v1.WorkflowsClient() + name = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + try: + operation = client.delete_workflow(name=name) + print("Waiting for the workflow to be deleted...") + operation.result(timeout=300) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to delete workflow '{workflow_id}' " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {name}") + print("State: DELETED") + + +# [END workflows_delete_workflow] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=delete_workflow.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + delete_workflow( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/execute_with_arguments.py b/workflows/cloud-client/execute_with_arguments.py index 0865cd321cc..6ddba256d06 100644 --- a/workflows/cloud-client/execute_with_arguments.py +++ b/workflows/cloud-client/execute_with_arguments.py @@ -4,7 +4,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -12,90 +12,91 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +# [START workflows_execute_with_arguments] +import json +import sys +import time -from google.cloud.workflows.executions_v1 import Execution +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 def execute_workflow_with_arguments( - project_id: str, - location: str, - workflow_id: str -) -> Execution: - """Execute a workflow and print the execution results. + project_id: str, location_id: str, workflow_id: str +) -> None: + """Runs a workflow with input arguments and waits for it to finish. Args: - project: The ID of the Google Cloud project - which contains the workflow to execute. - location: The location for the workflow. - workflow: The ID of the workflow to execute. - - Returns: - The execution response. + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow to run (for example, 'myFirstWorkflow'). """ + client = executions_v1.ExecutionsClient() + parent = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + # The workflow receives the argument as its runtime input (params: [args]). + argument = json.dumps({"searchTerm": "Cloud"}) + execution = executions_v1.Execution(argument=argument) + + try: + execution = client.create_execution(parent=parent, execution=execution) + print(f"Waiting for execution {execution.name} to finish...") + # Executions run asynchronously. Poll with exponential backoff until done. + delay_seconds = 1 + while execution.state == executions_v1.Execution.State.ACTIVE: + time.sleep(delay_seconds) + delay_seconds = min(delay_seconds * 2, 30) + execution = client.get_execution(name=execution.name) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to run workflow '{workflow_id}' in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {execution.name}") + if execution.state == executions_v1.Execution.State.FAILED: + print(f"Error: {execution.error.payload}") + print(f"Result: {execution.result}") + print(f"State: {execution.state.name}") -# [START workflows_execute_with_arguments] - import time - - from google.cloud import workflows_v1 - from google.cloud.workflows import executions_v1 - - from google.cloud.workflows.executions_v1.types import executions - # TODO(developer): Update and uncomment the following lines. - # project_id = "YOUR_PROJECT_ID" - # location = "YOUR_LOCATION" # For example: us-central1 - # workflow_id = "YOUR_WORKFLOW_ID" # For example: myFirstWorkflow +# [END workflows_execute_with_arguments] - # Initialize API clients. - execution_client = executions_v1.ExecutionsClient() - workflows_client = workflows_v1.WorkflowsClient() - # Construct the fully qualified location path. - parent = workflows_client.workflow_path(project_id, location, workflow_id) +if __name__ == "__main__": + import argparse + import os - # Execute the workflow adding an dictionary of arguments. - # Find more information about the Execution object here: - # https://cloud.google.com/python/docs/reference/workflows/latest/google.cloud.workflows.executions_v1.types.Execution - execution = executions_v1.Execution( - name=parent, - argument='{"searchTerm": "Cloud"}', + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=execute_workflow_with_arguments.__doc__.splitlines()[0] ) - - response = execution_client.create_execution( - parent=parent, - execution=execution, + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", ) - print(f"Created execution: {response.name}") - - # Wait for execution to finish, then print results. - execution_finished = False - backoff_delay = 1 # Start wait with delay of 1 second. - print("Poll for result...") - - # Keep polling the state until the execution finishes, - # using exponential backoff. - while not execution_finished: - execution = execution_client.get_execution( - request={"name": response.name} - ) - execution_finished = execution.state != executions.Execution.State.ACTIVE - - # If we haven't seen the result yet, keep waiting. - if not execution_finished: - print("- Waiting for results...") - time.sleep(backoff_delay) - # Double the delay to provide exponential backoff. - backoff_delay *= 2 - else: - print(f"Execution finished with state: {execution.state.name}") - print(f"Execution results: {execution.result}") -# [END workflows_execute_with_arguments] - return execution - - -if __name__ == "__main__": - PROJECT = os.getenv("GOOGLE_CLOUD_PROJECT") - assert PROJECT, "'GOOGLE_CLOUD_PROJECT' environment variable not set." + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() - execute_workflow_with_arguments(PROJECT, "us-central1", "myFirstWorkflow") + execute_workflow_with_arguments( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/execute_with_arguments_test.py b/workflows/cloud-client/execute_with_arguments_test.py deleted file mode 100644 index a465a4028f0..00000000000 --- a/workflows/cloud-client/execute_with_arguments_test.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import backoff - -from google.cloud.workflows.executions_v1.types import executions - -import execute_with_arguments - - -@backoff.on_exception(backoff.expo, AssertionError, max_tries=5) -def test_workflow_execution_with_arguments( - project_id: str, location: str, workflow_id: str -) -> None: - execution_result = execute_with_arguments.execute_workflow_with_arguments( - project_id, location, workflow_id - ) - assert execution_result.state == executions.Execution.State.SUCCEEDED - assert "searchTerm" in execution_result.argument - assert "Cloud" in execution_result.result diff --git a/workflows/cloud-client/execute_without_arguments.py b/workflows/cloud-client/execute_without_arguments.py index dd0980c9228..b68f3f69185 100644 --- a/workflows/cloud-client/execute_without_arguments.py +++ b/workflows/cloud-client/execute_without_arguments.py @@ -4,7 +4,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -12,83 +12,86 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +# [START workflows_execute_without_arguments] +import sys +import time -from google.cloud.workflows.executions_v1 import Execution +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 def execute_workflow_without_arguments( - project_id: str, - location: str, - workflow_id: str -) -> Execution: - """Execute a workflow and print the execution results. - - A workflow consists of a series of steps described - using the Workflows syntax, and can be written in either YAML or JSON. + project_id: str, location_id: str, workflow_id: str +) -> None: + """Runs a workflow and waits for it to finish. Args: - project: The ID of the Google Cloud project - which contains the workflow to execute. - location: The location for the workflow. - workflow: The ID of the workflow to execute. - - Returns: - The execution response. + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow to run (for example, 'myFirstWorkflow'). """ - -# [START workflows_execute_without_arguments] - import time - - from google.cloud import workflows_v1 - from google.cloud.workflows import executions_v1 - - from google.cloud.workflows.executions_v1.types import executions - - # TODO(developer): Update and uncomment the following lines. - # project_id = "YOUR_PROJECT_ID" - # location = "YOUR_LOCATION" # For example: us-central1 - # workflow_id = "YOUR_WORKFLOW_ID" # For example: myFirstWorkflow - - # Initialize API clients. - execution_client = executions_v1.ExecutionsClient() - workflows_client = workflows_v1.WorkflowsClient() - - # Construct the fully qualified location path. - parent = workflows_client.workflow_path(project_id, location, workflow_id) - - # Execute the workflow. - response = execution_client.create_execution(request={"parent": parent}) - print(f"Created execution: {response.name}") - - # Wait for execution to finish, then print results. - execution_finished = False - backoff_delay = 1 # Start wait with delay of 1 second. - print("Poll for result...") - - # Keep polling the state until the execution finishes, - # using exponential backoff. - while not execution_finished: - execution = execution_client.get_execution( - request={"name": response.name} + client = executions_v1.ExecutionsClient() + parent = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + try: + execution = client.create_execution(parent=parent) + print(f"Waiting for execution {execution.name} to finish...") + # Executions run asynchronously. Poll with exponential backoff until done. + delay_seconds = 1 + while execution.state == executions_v1.Execution.State.ACTIVE: + time.sleep(delay_seconds) + delay_seconds = min(delay_seconds * 2, 30) + execution = client.get_execution(name=execution.name) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to run workflow '{workflow_id}' in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, ) - execution_finished = execution.state != executions.Execution.State.ACTIVE + else: + print(f"Resource ID: {execution.name}") + if execution.state == executions_v1.Execution.State.FAILED: + print(f"Error: {execution.error.payload}") + print(f"Result: {execution.result}") + print(f"State: {execution.state.name}") + - # If we haven't seen the result yet, keep waiting. - if not execution_finished: - print("- Waiting for results...") - time.sleep(backoff_delay) - # Double the delay to provide exponential backoff. - backoff_delay *= 2 - else: - print(f"Execution finished with state: {execution.state.name}") - print(f"Execution results: {execution.result}") # [END workflows_execute_without_arguments] - return execution if __name__ == "__main__": - PROJECT = os.getenv("GOOGLE_CLOUD_PROJECT") - assert PROJECT, "'GOOGLE_CLOUD_PROJECT' environment variable not set." - - execute_workflow_without_arguments(PROJECT, "us-central1", "myFirstWorkflow") + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=execute_workflow_without_arguments.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + execute_workflow_without_arguments( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/execute_without_arguments_test.py b/workflows/cloud-client/execute_without_arguments_test.py deleted file mode 100644 index c8ff8050b9e..00000000000 --- a/workflows/cloud-client/execute_without_arguments_test.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright 2021 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import backoff - -from google.cloud.workflows.executions_v1.types import executions - -import execute_without_arguments - - -@backoff.on_exception(backoff.expo, AssertionError, max_tries=5) -def test_workflow_execution_without_arguments( - project_id: str, location: str, workflow_id: str -) -> None: - result = execute_without_arguments.execute_workflow_without_arguments( - project_id, location, workflow_id - ) - assert result.state == executions.Execution.State.SUCCEEDED - assert result.result diff --git a/workflows/cloud-client/get_execution.py b/workflows/cloud-client/get_execution.py new file mode 100644 index 00000000000..eed2a23f84b --- /dev/null +++ b/workflows/cloud-client/get_execution.py @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_get_execution] +import sys + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 + + +def get_execution( + project_id: str, location_id: str, workflow_id: str, execution_id: str +) -> None: + """Gets an execution, including its state and its result or error. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'myFirstWorkflow'). + execution_id: The ID of the execution. + """ + client = executions_v1.ExecutionsClient() + name = ( + f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + f"/executions/{execution_id}" + ) + + try: + execution = client.get_execution(name=name) + except exceptions.NotFound as err: + print( + f"error: execution '{execution_id}' not found " + f"for workflow '{workflow_id}' in location '{location_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to get execution '{execution_id}' " + f"for workflow '{workflow_id}' in location '{location_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {execution.name}") + # The result is set when the execution succeeds; the error when it fails. + print(f"Error: {execution.error.payload}") + print(f"Result: {execution.result}") + print(f"State: {execution.state.name}") + + +# [END workflows_get_execution] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser(description=get_execution.__doc__.splitlines()[0]) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + parser.add_argument( + "--execution-id", + default=os.environ.get("EXECUTION_ID", "example-execution-id"), + help="ID of the execution (default: $EXECUTION_ID)", + ) + args = parser.parse_args() + + get_execution( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + execution_id=args.execution_id, + ) diff --git a/workflows/cloud-client/get_workflow.py b/workflows/cloud-client/get_workflow.py new file mode 100644 index 00000000000..ead0551be62 --- /dev/null +++ b/workflows/cloud-client/get_workflow.py @@ -0,0 +1,85 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_get_workflow] +import sys + +from google.api_core import exceptions +from google.cloud import workflows_v1 + + +def get_workflow(project_id: str, location_id: str, workflow_id: str) -> None: + """Gets a workflow, including its current revision and state. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'example-workflow'). + """ + client = workflows_v1.WorkflowsClient() + name = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + try: + workflow = client.get_workflow(name=name) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to get workflow '{workflow_id}' " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {workflow.name}") + print(f"Creation Date: {workflow.create_time}") + print(f"Revision ID: {workflow.revision_id}") + print(f"Service Account: {workflow.service_account}") + print(f"State: {workflow.state.name}") + + +# [END workflows_get_workflow] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser(description=get_workflow.__doc__.splitlines()[0]) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + get_workflow( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/list_executions.py b/workflows/cloud-client/list_executions.py new file mode 100644 index 00000000000..f590c04535b --- /dev/null +++ b/workflows/cloud-client/list_executions.py @@ -0,0 +1,92 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_list_executions] +import sys + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 + + +def list_executions(project_id: str, location_id: str, workflow_id: str) -> None: + """Lists a workflow's failed executions, newest first. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'myFirstWorkflow'). + """ + client = executions_v1.ExecutionsClient() + parent = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + request = executions_v1.ListExecutionsRequest( + parent=parent, + filter='state="FAILED"', + order_by="createTime desc", + # The default BASIC view leaves out each execution's error and result. + view=executions_v1.ExecutionView.FULL, + ) + + try: + # The client fetches more pages as the loop needs them. + for execution in client.list_executions(request=request): + print(f"Resource ID: {execution.name}") + print(f"Error: {execution.error.payload}") + print(f"Start Time: {execution.start_time}") + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to list executions " + f"for workflow '{workflow_id}' in location '{location_id}': {err.message}", + file=sys.stderr, + ) + + +# [END workflows_list_executions] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=list_executions.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + list_executions( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/list_workflows.py b/workflows/cloud-client/list_workflows.py new file mode 100644 index 00000000000..524ec2222e3 --- /dev/null +++ b/workflows/cloud-client/list_workflows.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_client_library] +import sys + +from google.api_core import exceptions +from google.cloud import workflows_v1 + + +def list_workflows(project_id: str, location_id: str) -> None: + """Lists the workflows in a location. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + """ + client = workflows_v1.WorkflowsClient() + parent = f"projects/{project_id}/locations/{location_id}" + + try: + # The client fetches more pages as the loop needs them. + for workflow in client.list_workflows(parent=parent): + print(f"Resource ID: {workflow.name}") + print(f"Creation Date: {workflow.create_time}") + print(f"Revision ID: {workflow.revision_id}") + print(f"State: {workflow.state.name}") + except exceptions.InvalidArgument as err: + print( + f"error: invalid location '{location_id}' " + f"for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to list workflows " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END workflows_client_library] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser(description=list_workflows.__doc__.splitlines()[0]) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + args = parser.parse_args() + + list_workflows( + project_id=args.project_id, + location_id=args.location_id, + ) diff --git a/workflows/cloud-client/main.py b/workflows/cloud-client/main.py index 4300df16d30..9ffeb32ad96 100644 --- a/workflows/cloud-client/main.py +++ b/workflows/cloud-client/main.py @@ -4,7 +4,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -12,87 +12,84 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os - -from google.cloud.workflows.executions_v1 import Execution +# [START workflows_api_quickstart] +import sys +import time +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 -def execute_workflow( - project_id: str, - location: str, - workflow_id: str -) -> Execution: - """Execute a workflow and print the execution results. - A workflow consists of a series of steps described - using the Workflows syntax, and can be written in either YAML or JSON. +def execute_workflow(project_id: str, location_id: str, workflow_id: str) -> None: + """Runs a workflow and waits for it to finish. Args: - project: The ID of the Google Cloud project - which contains the workflow to execute. - location: The location for the workflow. - workflow: The ID of the workflow to execute. - - Returns: - The execution response. + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow to run (for example, 'myFirstWorkflow'). """ - -# [START workflows_execute_without_arguments] -# [START workflows_api_quickstart] - import time - - from google.cloud import workflows_v1 - from google.cloud.workflows import executions_v1 - - from google.cloud.workflows.executions_v1.types import executions - - # TODO(developer): Update and uncomment the following lines. - # project_id = "MY_PROJECT_ID" - # location = "MY_LOCATION" # For example: us-central1 - # workflow_id = "MY_WORKFLOW_ID" # For example: myFirstWorkflow - - # [START workflows_api_quickstart_client_libraries] - # Initialize API clients. - execution_client = executions_v1.ExecutionsClient() - workflows_client = workflows_v1.WorkflowsClient() - # [END workflows_api_quickstart_client_libraries] - - # [START workflows_api_quickstart_execution] - # Construct the fully qualified location path. - parent = workflows_client.workflow_path(project_id, location, workflow_id) - - # Execute the workflow. - response = execution_client.create_execution(request={"parent": parent}) - print(f"Created execution: {response.name}") - - # Wait for execution to finish, then print results. - execution_finished = False - backoff_delay = 1 # Start wait with delay of 1 second. - print("Poll for result...") - - while not execution_finished: - execution = execution_client.get_execution( - request={"name": response.name} + client = executions_v1.ExecutionsClient() + parent = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + try: + execution = client.create_execution(parent=parent) + print(f"Waiting for execution {execution.name} to finish...") + # Executions run asynchronously. Poll with exponential backoff until done. + delay_seconds = 1 + while execution.state == executions_v1.Execution.State.ACTIVE: + time.sleep(delay_seconds) + delay_seconds = min(delay_seconds * 2, 30) + execution = client.get_execution(name=execution.name) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to run workflow '{workflow_id}' in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, ) - execution_finished = execution.state != executions.Execution.State.ACTIVE + else: + print(f"Resource ID: {execution.name}") + if execution.state == executions_v1.Execution.State.FAILED: + print(f"Error: {execution.error.payload}") + print(f"Result: {execution.result}") + print(f"State: {execution.state.name}") + - # If we haven't seen the result yet, keep waiting. - if not execution_finished: - print("- Waiting for results...") - time.sleep(backoff_delay) - # Double the delay to provide exponential backoff. - backoff_delay *= 2 - else: - print(f"Execution finished with state: {execution.state.name}") - print(f"Execution results: {execution.result}") - # [END workflows_api_quickstart_execution] # [END workflows_api_quickstart] -# [END workflows_execute_without_arguments] - return execution if __name__ == "__main__": - PROJECT = os.getenv("GOOGLE_CLOUD_PROJECT") - assert PROJECT, "'GOOGLE_CLOUD_PROJECT' environment variable not set." - - execute_workflow(PROJECT, "us-central1", "myFirstWorkflow") + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=execute_workflow.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + execute_workflow( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/cloud-client/noxfile_config.py b/workflows/cloud-client/noxfile_config.py index 1b9d9e4a580..0dae3edd19a 100644 --- a/workflows/cloud-client/noxfile_config.py +++ b/workflows/cloud-client/noxfile_config.py @@ -22,7 +22,7 @@ TEST_CONFIG_OVERRIDE = { # You can opt out from the test for specific Python versions. - "ignored_versions": ["3.8", "3.9", "3.11", "3.12", "3.13"], + "ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"], # Old samples are opted out of enforcing Python type hints # All new samples should feature them "enforce_type_hints": True, diff --git a/workflows/cloud-client/requirements-test.txt b/workflows/cloud-client/requirements-test.txt index 975d5ee58c2..b3b2b3ca4cc 100644 --- a/workflows/cloud-client/requirements-test.txt +++ b/workflows/cloud-client/requirements-test.txt @@ -1,2 +1 @@ -pytest==9.0.3; python_version >= "3.10" -backoff==2.2.1 +pytest==9.0.3 diff --git a/workflows/cloud-client/requirements.txt b/workflows/cloud-client/requirements.txt index e9f5097f4ed..ee1960cac8c 100644 --- a/workflows/cloud-client/requirements.txt +++ b/workflows/cloud-client/requirements.txt @@ -1 +1 @@ -google-cloud-workflows==1.18.2 +google-cloud-workflows==1.24.0 diff --git a/workflows/cloud-client/tests/cancel_execution_test.py b/workflows/cloud-client/tests/cancel_execution_test.py new file mode 100644 index 00000000000..a06d516a760 --- /dev/null +++ b/workflows/cloud-client/tests/cancel_execution_test.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import cancel_execution # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(EXECUTIONS_CLIENT) +def test_cancel_execution_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.cancel_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.CANCELLED + ) + + cancel_execution.cancel_execution( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID, EXECUTION_ID + ) + + assert client.cancel_execution.call_args.kwargs["name"] == EXECUTION_NAME + out, err = capsys.readouterr() + assert "State: CANCELLED" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: execution 'test-execution' not found", + ), + ( + exceptions.PermissionDenied("test detail"), + "error: failed to cancel execution", + ), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_cancel_execution_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.cancel_execution.side_effect = error + + cancel_execution.cancel_execution( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID, EXECUTION_ID + ) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/main_test.py b/workflows/cloud-client/tests/conftest.py similarity index 50% rename from workflows/cloud-client/main_test.py rename to workflows/cloud-client/tests/conftest.py index 0ca8d202db9..9f7ff1b1c8e 100644 --- a/workflows/cloud-client/main_test.py +++ b/workflows/cloud-client/tests/conftest.py @@ -1,10 +1,10 @@ -# Copyright 2021 Google LLC +# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -12,15 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import backoff +import pathlib +import sys -from google.cloud.workflows.executions_v1.types import executions - -import main - - -@backoff.on_exception(backoff.expo, AssertionError, max_tries=5) -def test_workflow_execution(project_id: str, location: str, workflow_id: str) -> None: - result = main.execute_workflow(project_id, location, workflow_id) - assert result.state == executions.Execution.State.SUCCEEDED - assert result.result +# The samples live one directory up. The directory name (cloud-client) isn't a +# valid package name, so add it to the import path instead. +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) diff --git a/workflows/cloud-client/tests/create_workflow_test.py b/workflows/cloud-client/tests/create_workflow_test.py new file mode 100644 index 00000000000..f39a62e5581 --- /dev/null +++ b/workflows/cloud-client/tests/create_workflow_test.py @@ -0,0 +1,102 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import workflows_v1 +import pytest + +import create_workflow # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(WORKFLOWS_CLIENT) +def test_create_workflow_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.create_workflow.return_value.result.return_value = workflows_v1.Workflow( + name=WORKFLOW_NAME, + revision_id="000002-abc", + state=workflows_v1.Workflow.State.ACTIVE, + description="example-description", + labels={"environment": "example-updated"}, + call_log_level=workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY, + ) + + create_workflow.create_workflow( + PROJECT_ID, + LOCATION_ID, + WORKFLOW_ID, + "runner@test-project.iam.gserviceaccount.com", + ) + + kwargs = client.create_workflow.call_args.kwargs + assert kwargs["parent"] == PARENT + assert kwargs["workflow_id"] == WORKFLOW_ID + workflow = kwargs["workflow"] + assert workflow.service_account == "runner@test-project.iam.gserviceaccount.com" + assert workflow.call_log_level == workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY + assert workflow.user_env_vars["EXAMPLE_GREETING"] == "example-greeting" + assert "EXAMPLE_GREETING" in workflow.source_contents + out, err = capsys.readouterr() + assert "Waiting for the workflow to deploy..." in out + assert f"Resource ID: {WORKFLOW_NAME}" in out + assert "Call Log Level: LOG_ERRORS_ONLY" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.AlreadyExists("test detail"), + "error: workflow 'test-workflow' already exists", + ), + ( + exceptions.PermissionDenied("test detail"), + "error: failed to create workflow", + ), + ], +) +@mock.patch(WORKFLOWS_CLIENT) +def test_create_workflow_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.create_workflow.side_effect = error + + create_workflow.create_workflow( + PROJECT_ID, + LOCATION_ID, + WORKFLOW_ID, + "runner@test-project.iam.gserviceaccount.com", + ) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/delete_workflow_test.py b/workflows/cloud-client/tests/delete_workflow_test.py new file mode 100644 index 00000000000..8ca0ee31535 --- /dev/null +++ b/workflows/cloud-client/tests/delete_workflow_test.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +import pytest + +import delete_workflow # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(WORKFLOWS_CLIENT) +def test_delete_workflow_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + pass + + delete_workflow.delete_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + assert client.delete_workflow.call_args.kwargs["name"] == WORKFLOW_NAME + client.delete_workflow.return_value.result.assert_called_once() + out, err = capsys.readouterr() + assert f"Resource ID: {WORKFLOW_NAME}" in out + assert "State: DELETED" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + ( + exceptions.PermissionDenied("test detail"), + "error: failed to delete workflow", + ), + ], +) +@mock.patch(WORKFLOWS_CLIENT) +def test_delete_workflow_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.delete_workflow.side_effect = error + + delete_workflow.delete_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/entrypoint_test.py b/workflows/cloud-client/tests/entrypoint_test.py new file mode 100644 index 00000000000..f769b30b6b9 --- /dev/null +++ b/workflows/cloud-client/tests/entrypoint_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pathlib +import runpy +from unittest import mock + +from google.cloud import workflows_v1 +import pytest + +SAMPLE = str(pathlib.Path(__file__).resolve().parent.parent / "get_workflow.py") +CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +ENV_VARS = ("GOOGLE_CLOUD_PROJECT", "GOOGLE_CLOUD_LOCATION", "WORKFLOW_ID") + + +def run_entrypoint(monkeypatch: pytest.MonkeyPatch, *flags: str) -> str: + """Runs get_workflow.py as a script and returns the workflow it requested.""" + monkeypatch.setattr("sys.argv", ["get_workflow.py", *flags]) + with mock.patch(CLIENT) as mock_client_class: + client = mock_client_class.return_value + client.get_workflow.return_value = workflows_v1.Workflow(name="workflow") + runpy.run_path(SAMPLE, run_name="__main__") + return client.get_workflow.call_args.kwargs["name"] + + +def test_entrypoint_reads_environment_variables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for env_var, value in zip(ENV_VARS, ("env-project", "europe-west1", "env-wf")): + monkeypatch.setenv(env_var, value) + + name = run_entrypoint(monkeypatch) + + assert name == "projects/env-project/locations/europe-west1/workflows/env-wf" + + +def test_entrypoint_flags_override_environment_variables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for env_var, value in zip(ENV_VARS, ("env-project", "europe-west1", "env-wf")): + monkeypatch.setenv(env_var, value) + + name = run_entrypoint(monkeypatch, "--workflow-id", "flag-wf") + + assert name == "projects/env-project/locations/europe-west1/workflows/flag-wf" + + +def test_entrypoint_uses_example_defaults(monkeypatch: pytest.MonkeyPatch) -> None: + for env_var in ENV_VARS: + monkeypatch.delenv(env_var, raising=False) + + name = run_entrypoint(monkeypatch) + + assert name == ( + "projects/example-project-id/locations/us-central1/workflows/myFirstWorkflow" + ) diff --git a/workflows/cloud-client/tests/execute_with_arguments_test.py b/workflows/cloud-client/tests/execute_with_arguments_test.py new file mode 100644 index 00000000000..af75da23d3a --- /dev/null +++ b/workflows/cloud-client/tests/execute_with_arguments_test.py @@ -0,0 +1,120 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import execute_with_arguments # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch.object(execute_with_arguments.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_with_arguments_waits_for_result( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ) + client.get_execution.side_effect = [ + executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ), + executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.SUCCEEDED, + result='"test result"', + ), + ] + + execute_with_arguments.execute_workflow_with_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + assert client.create_execution.call_args.kwargs["parent"] == WORKFLOW_NAME + execution = client.create_execution.call_args.kwargs["execution"] + assert execution.argument == '{"searchTerm": "Cloud"}' + assert mock_sleep.call_count == 2 + out, err = capsys.readouterr() + assert f"Resource ID: {EXECUTION_NAME}" in out + assert 'Result: "test result"' in out + assert "State: SUCCEEDED" in out + assert err == "" + + +@mock.patch.object(execute_with_arguments.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_with_arguments_prints_error_when_execution_fails( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.FAILED, + error=executions_v1.Execution.Error(payload="test failure"), + ) + + execute_with_arguments.execute_workflow_with_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + out, _ = capsys.readouterr() + assert "Error: test failure" in out + assert "State: FAILED" in out + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to run workflow"), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_with_arguments_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.create_execution.side_effect = error + + execute_with_arguments.execute_workflow_with_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/execute_without_arguments_test.py b/workflows/cloud-client/tests/execute_without_arguments_test.py new file mode 100644 index 00000000000..f1d66333edd --- /dev/null +++ b/workflows/cloud-client/tests/execute_without_arguments_test.py @@ -0,0 +1,119 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import execute_without_arguments # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch.object(execute_without_arguments.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_without_arguments_waits_for_result( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ) + client.get_execution.side_effect = [ + executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ), + executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.SUCCEEDED, + result='"test result"', + ), + ] + + execute_without_arguments.execute_workflow_without_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + assert client.create_execution.call_args.kwargs["parent"] == WORKFLOW_NAME + + assert mock_sleep.call_count == 2 + out, err = capsys.readouterr() + assert f"Resource ID: {EXECUTION_NAME}" in out + assert 'Result: "test result"' in out + assert "State: SUCCEEDED" in out + assert err == "" + + +@mock.patch.object(execute_without_arguments.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_without_arguments_prints_error_when_execution_fails( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.FAILED, + error=executions_v1.Execution.Error(payload="test failure"), + ) + + execute_without_arguments.execute_workflow_without_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + out, _ = capsys.readouterr() + assert "Error: test failure" in out + assert "State: FAILED" in out + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to run workflow"), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_without_arguments_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.create_execution.side_effect = error + + execute_without_arguments.execute_workflow_without_arguments( + PROJECT_ID, LOCATION_ID, WORKFLOW_ID + ) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/get_execution_test.py b/workflows/cloud-client/tests/get_execution_test.py new file mode 100644 index 00000000000..7100f1a5615 --- /dev/null +++ b/workflows/cloud-client/tests/get_execution_test.py @@ -0,0 +1,78 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import get_execution # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(EXECUTIONS_CLIENT) +def test_get_execution_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.get_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.SUCCEEDED, + result='"done"', + ) + + get_execution.get_execution(PROJECT_ID, LOCATION_ID, WORKFLOW_ID, EXECUTION_ID) + + assert client.get_execution.call_args.kwargs["name"] == EXECUTION_NAME + out, err = capsys.readouterr() + assert 'Result: "done"' in out + assert "State: SUCCEEDED" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: execution 'test-execution' not found", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to get execution"), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_get_execution_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.get_execution.side_effect = error + + get_execution.get_execution(PROJECT_ID, LOCATION_ID, WORKFLOW_ID, EXECUTION_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/get_workflow_test.py b/workflows/cloud-client/tests/get_workflow_test.py new file mode 100644 index 00000000000..7e28c16f461 --- /dev/null +++ b/workflows/cloud-client/tests/get_workflow_test.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import workflows_v1 +import pytest + +import get_workflow # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(WORKFLOWS_CLIENT) +def test_get_workflow_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.get_workflow.return_value = workflows_v1.Workflow( + name=WORKFLOW_NAME, + revision_id="000002-abc", + state=workflows_v1.Workflow.State.ACTIVE, + description="example-description", + labels={"environment": "example-updated"}, + call_log_level=workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY, + ) + + get_workflow.get_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + assert client.get_workflow.call_args.kwargs["name"] == WORKFLOW_NAME + out, err = capsys.readouterr() + assert f"Resource ID: {WORKFLOW_NAME}" in out + assert "State: ACTIVE" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to get workflow"), + ], +) +@mock.patch(WORKFLOWS_CLIENT) +def test_get_workflow_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.get_workflow.side_effect = error + + get_workflow.get_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/integration_test.py b/workflows/cloud-client/tests/integration_test.py new file mode 100644 index 00000000000..514ef157d7f --- /dev/null +++ b/workflows/cloud-client/tests/integration_test.py @@ -0,0 +1,265 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests that run the Workflows samples live. + +These tests are skipped unless RUN_INTEGRATION_TESTS is "true". They need: + GOOGLE_CLOUD_PROJECT: Project that runs the tests. + SERVICE_ACCOUNT_EMAIL: Service account that test workflows run as. + GOOGLE_CLOUD_LOCATION: Region for the workflows (default 'us-central1'). +""" + +import os +import pathlib +import time +from typing import Iterator +import uuid + +from google.api_core import exceptions +from google.cloud import workflows_v1 +from google.cloud.workflows import executions_v1 +import pytest + +import cancel_execution # noqa: I100 +import create_workflow # noqa: I100 +import delete_workflow # noqa: I100 +import execute_with_arguments # noqa: I100 +import execute_without_arguments # noqa: I100 +import get_execution # noqa: I100 +import get_workflow # noqa: I100 +import list_executions # noqa: I100 +import list_workflows # noqa: I100 +import main # noqa: I100 +import update_workflow # noqa: I100 + +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_INTEGRATION_TESTS") != "true", + reason="Set RUN_INTEGRATION_TESTS=true to run integration tests.", +) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "") +LOCATION_ID = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") +SERVICE_ACCOUNT_EMAIL = os.environ.get("SERVICE_ACCOUNT_EMAIL", "") +MY_FIRST_WORKFLOW = ( + pathlib.Path(__file__).resolve().parent.parent / "myFirstWorkflow.workflows.yaml" +).read_text() +SLEEP_WORKFLOW = """ +main: + steps: + - wait: + call: sys.sleep + args: + seconds: 300 +""" +FAILING_WORKFLOW = """ +main: + steps: + - fail: + raise: "example-failure" +""" + + +@pytest.fixture(scope="module") +def workflows_client() -> workflows_v1.WorkflowsClient: + return workflows_v1.WorkflowsClient() + + +@pytest.fixture(scope="module") +def executions_client() -> executions_v1.ExecutionsClient: + return executions_v1.ExecutionsClient() + + +def deploy(client: workflows_v1.WorkflowsClient, source: str) -> Iterator[str]: + """Deploys a test workflow with a unique ID and deletes it afterward.""" + workflow_id = f"test-workflow-{uuid.uuid4().hex[:12]}" + client.create_workflow( + parent=f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + workflow=workflows_v1.Workflow( + source_contents=source, service_account=SERVICE_ACCOUNT_EMAIL + ), + workflow_id=workflow_id, + ).result(timeout=300) + yield workflow_id + try: + client.delete_workflow( + name=f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{workflow_id}" + ).result(timeout=300) + except exceptions.NotFound: + pass + + +@pytest.fixture(scope="module") +def my_first_workflow_id( + workflows_client: workflows_v1.WorkflowsClient, +) -> Iterator[str]: + yield from deploy(workflows_client, MY_FIRST_WORKFLOW) + + +@pytest.fixture +def sleep_workflow_id(workflows_client: workflows_v1.WorkflowsClient) -> Iterator[str]: + yield from deploy(workflows_client, SLEEP_WORKFLOW) + + +@pytest.fixture +def failing_workflow_id( + workflows_client: workflows_v1.WorkflowsClient, +) -> Iterator[str]: + yield from deploy(workflows_client, FAILING_WORKFLOW) + + +@pytest.fixture +def new_workflow_id(workflows_client: workflows_v1.WorkflowsClient) -> Iterator[str]: + """Provides an unused workflow ID and deletes the workflow afterward.""" + workflow_id = f"test-workflow-{uuid.uuid4().hex[:12]}" + yield workflow_id + try: + workflows_client.delete_workflow( + name=f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{workflow_id}" + ).result(timeout=300) + except exceptions.NotFound: + pass + + +def test_workflow_lifecycle( + capsys: pytest.CaptureFixture[str], new_workflow_id: str +) -> None: + args = (PROJECT_ID, LOCATION_ID, new_workflow_id) + name = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{new_workflow_id}" + + create_workflow.create_workflow(*args, SERVICE_ACCOUNT_EMAIL) + out, err = capsys.readouterr() + assert f"Resource ID: {name}" in out + assert "Call Log Level: LOG_ERRORS_ONLY" in out + assert "error:" not in err + + get_workflow.get_workflow(*args) + assert f"Service Account: projects/{PROJECT_ID}/serviceAccounts/" in ( + capsys.readouterr().out + ) + + update_workflow.update_workflow(*args) + assert "Description: example-description" in capsys.readouterr().out + + list_workflows.list_workflows(PROJECT_ID, LOCATION_ID) + assert f"Resource ID: {name}" in capsys.readouterr().out + + delete_workflow.delete_workflow(*args) + assert "State: DELETED" in capsys.readouterr().out + + get_workflow.get_workflow(*args) + assert f"error: workflow '{new_workflow_id}' not found" in capsys.readouterr().err + + +def test_created_workflow_reads_its_environment_variable( + capsys: pytest.CaptureFixture[str], new_workflow_id: str +) -> None: + create_workflow.create_workflow( + PROJECT_ID, LOCATION_ID, new_workflow_id, SERVICE_ACCOUNT_EMAIL + ) + capsys.readouterr() + + execute_without_arguments.execute_workflow_without_arguments( + PROJECT_ID, LOCATION_ID, new_workflow_id + ) + + out = capsys.readouterr().out + assert 'Result: "example-greeting"' in out + assert "State: SUCCEEDED" in out + + +@pytest.mark.parametrize( + "run", + [ + execute_without_arguments.execute_workflow_without_arguments, + main.execute_workflow, + ], +) +def test_execute_my_first_workflow( + capsys: pytest.CaptureFixture[str], my_first_workflow_id: str, run: object +) -> None: + run(PROJECT_ID, LOCATION_ID, my_first_workflow_id) + + out, err = capsys.readouterr() + assert "State: SUCCEEDED" in out + assert "error:" not in err + + +def test_execute_with_arguments( + capsys: pytest.CaptureFixture[str], my_first_workflow_id: str +) -> None: + execute_with_arguments.execute_workflow_with_arguments( + PROJECT_ID, LOCATION_ID, my_first_workflow_id + ) + + out = capsys.readouterr().out + assert "State: SUCCEEDED" in out + assert "Cloud" in out + + +def test_get_execution( + capsys: pytest.CaptureFixture[str], + executions_client: executions_v1.ExecutionsClient, + my_first_workflow_id: str, +) -> None: + parent = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{my_first_workflow_id}" + execution = executions_client.create_execution(parent=parent) + execution_id = execution.name.rsplit("/", 1)[-1] + + get_execution.get_execution( + PROJECT_ID, LOCATION_ID, my_first_workflow_id, execution_id + ) + + out, err = capsys.readouterr() + assert f"Resource ID: {execution.name}" in out + assert "error:" not in err + + +def test_list_failed_executions( + capsys: pytest.CaptureFixture[str], + executions_client: executions_v1.ExecutionsClient, + failing_workflow_id: str, +) -> None: + parent = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{failing_workflow_id}" + ) + execution = executions_client.create_execution(parent=parent) + deadline = time.monotonic() + 120 + while execution.state == executions_v1.Execution.State.ACTIVE: + assert time.monotonic() < deadline, "The execution didn't finish." + time.sleep(2) + execution = executions_client.get_execution(name=execution.name) + + list_executions.list_executions(PROJECT_ID, LOCATION_ID, failing_workflow_id) + + out = capsys.readouterr().out + assert f"Resource ID: {execution.name}" in out + assert "example-failure" in out + + +def test_cancel_execution( + capsys: pytest.CaptureFixture[str], + executions_client: executions_v1.ExecutionsClient, + sleep_workflow_id: str, +) -> None: + parent = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{sleep_workflow_id}" + ) + execution = executions_client.create_execution(parent=parent) + execution_id = execution.name.rsplit("/", 1)[-1] + + cancel_execution.cancel_execution( + PROJECT_ID, LOCATION_ID, sleep_workflow_id, execution_id + ) + + assert "State: CANCELLED" in capsys.readouterr().out diff --git a/workflows/cloud-client/tests/list_executions_test.py b/workflows/cloud-client/tests/list_executions_test.py new file mode 100644 index 00000000000..f697495e422 --- /dev/null +++ b/workflows/cloud-client/tests/list_executions_test.py @@ -0,0 +1,87 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import list_executions # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(EXECUTIONS_CLIENT) +def test_list_executions_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.list_executions.return_value = [ + executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.FAILED, + error=executions_v1.Execution.Error(payload="boom"), + ) + ] + + list_executions.list_executions(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + request = client.list_executions.call_args.kwargs["request"] + assert request.parent == WORKFLOW_NAME + assert request.filter == 'state="FAILED"' + assert request.order_by == "createTime desc" + assert request.view == executions_v1.ExecutionView.FULL + out, err = capsys.readouterr() + assert f"Resource ID: {EXECUTION_NAME}" in out + assert "Error: boom" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + ( + exceptions.PermissionDenied("test detail"), + "error: failed to list executions", + ), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_list_executions_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.list_executions.side_effect = error + + list_executions.list_executions(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/list_workflows_test.py b/workflows/cloud-client/tests/list_workflows_test.py new file mode 100644 index 00000000000..7cea8300a1c --- /dev/null +++ b/workflows/cloud-client/tests/list_workflows_test.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import workflows_v1 +import pytest + +import list_workflows # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(WORKFLOWS_CLIENT) +def test_list_workflows_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.list_workflows.return_value = [ + workflows_v1.Workflow( + name=WORKFLOW_NAME, + revision_id="000002-abc", + state=workflows_v1.Workflow.State.ACTIVE, + description="example-description", + labels={"environment": "example-updated"}, + call_log_level=workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY, + ) + ] + + list_workflows.list_workflows(PROJECT_ID, LOCATION_ID) + + assert client.list_workflows.call_args.kwargs["parent"] == PARENT + out, err = capsys.readouterr() + assert f"Resource ID: {WORKFLOW_NAME}" in out + assert "Revision ID: 000002-abc" in out + assert "State: ACTIVE" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.InvalidArgument("test detail"), + "error: invalid location 'us-central1'", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to list workflows"), + ], +) +@mock.patch(WORKFLOWS_CLIENT) +def test_list_workflows_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.list_workflows.side_effect = error + + list_workflows.list_workflows(PROJECT_ID, LOCATION_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/main_test.py b/workflows/cloud-client/tests/main_test.py new file mode 100644 index 00000000000..8a4096ecbf0 --- /dev/null +++ b/workflows/cloud-client/tests/main_test.py @@ -0,0 +1,113 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +import main # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch.object(main.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_waits_for_result( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ) + client.get_execution.side_effect = [ + executions_v1.Execution( + name=EXECUTION_NAME, state=executions_v1.Execution.State.ACTIVE + ), + executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.SUCCEEDED, + result='"test result"', + ), + ] + + main.execute_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + assert client.create_execution.call_args.kwargs["parent"] == WORKFLOW_NAME + + assert mock_sleep.call_count == 2 + out, err = capsys.readouterr() + assert f"Resource ID: {EXECUTION_NAME}" in out + assert 'Result: "test result"' in out + assert "State: SUCCEEDED" in out + assert err == "" + + +@mock.patch.object(main.time, "sleep") +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_prints_error_when_execution_fails( + mock_client_class: mock.MagicMock, + mock_sleep: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + client = mock_client_class.return_value + client.create_execution.return_value = executions_v1.Execution( + name=EXECUTION_NAME, + state=executions_v1.Execution.State.FAILED, + error=executions_v1.Execution.Error(payload="test failure"), + ) + + main.execute_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, _ = capsys.readouterr() + assert "Error: test failure" in out + assert "State: FAILED" in out + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + (exceptions.PermissionDenied("test detail"), "error: failed to run workflow"), + ], +) +@mock.patch(EXECUTIONS_CLIENT) +def test_execute_workflow_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.create_execution.side_effect = error + + main.execute_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/tests/update_workflow_test.py b/workflows/cloud-client/tests/update_workflow_test.py new file mode 100644 index 00000000000..3802ba86447 --- /dev/null +++ b/workflows/cloud-client/tests/update_workflow_test.py @@ -0,0 +1,86 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import workflows_v1 +import pytest + +import update_workflow # noqa: I100 + +PROJECT_ID = "test-project" +LOCATION_ID = "us-central1" +WORKFLOW_ID = "test-workflow" +EXECUTION_ID = "test-execution" +PARENT = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" +WORKFLOW_NAME = f"{PARENT}/workflows/{WORKFLOW_ID}" +EXECUTION_NAME = f"{WORKFLOW_NAME}/executions/{EXECUTION_ID}" +WORKFLOWS_CLIENT = "google.cloud.workflows_v1.WorkflowsClient" +EXECUTIONS_CLIENT = "google.cloud.workflows.executions_v1.ExecutionsClient" + + +@mock.patch(WORKFLOWS_CLIENT) +def test_update_workflow_success( + mock_client_class: mock.MagicMock, capsys: pytest.CaptureFixture[str] +) -> None: + client = mock_client_class.return_value + client.update_workflow.return_value.result.return_value = workflows_v1.Workflow( + name=WORKFLOW_NAME, + revision_id="000002-abc", + state=workflows_v1.Workflow.State.ACTIVE, + description="example-description", + labels={"environment": "example-updated"}, + call_log_level=workflows_v1.Workflow.CallLogLevel.LOG_ERRORS_ONLY, + ) + + update_workflow.update_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + kwargs = client.update_workflow.call_args.kwargs + assert kwargs["workflow"].name == WORKFLOW_NAME + assert list(kwargs["update_mask"].paths) == ["description", "labels"] + out, err = capsys.readouterr() + assert "Description: example-description" in out + assert "Labels: {'environment': 'example-updated'}" in out + assert err == "" + + +@pytest.mark.parametrize( + "error, expected", + [ + ( + exceptions.NotFound("test detail"), + "error: workflow 'test-workflow' not found", + ), + ( + exceptions.PermissionDenied("test detail"), + "error: failed to update workflow", + ), + ], +) +@mock.patch(WORKFLOWS_CLIENT) +def test_update_workflow_errors( + mock_client_class: mock.MagicMock, + capsys: pytest.CaptureFixture[str], + error: Exception, + expected: str, +) -> None: + mock_client_class.return_value.update_workflow.side_effect = error + + update_workflow.update_workflow(PROJECT_ID, LOCATION_ID, WORKFLOW_ID) + + out, err = capsys.readouterr() + assert out == "" + assert err.startswith(expected) + assert "test detail" in err diff --git a/workflows/cloud-client/update_workflow.py b/workflows/cloud-client/update_workflow.py new file mode 100644 index 00000000000..63df55d54f8 --- /dev/null +++ b/workflows/cloud-client/update_workflow.py @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_update_workflow] +import sys + +from google.api_core import exceptions +from google.cloud import workflows_v1 +from google.protobuf import field_mask_pb2 + + +def update_workflow(project_id: str, location_id: str, workflow_id: str) -> None: + """Updates a workflow's description and labels. + + Args: + project_id: The Google Cloud project ID (for example, 'example-project-id'). + location_id: The region of the workflow (for example, 'us-central1'). + workflow_id: The ID of the workflow (for example, 'example-workflow'). + """ + client = workflows_v1.WorkflowsClient() + name = f"projects/{project_id}/locations/{location_id}/workflows/{workflow_id}" + + workflow = workflows_v1.Workflow( + name=name, + description="example-description", + labels={"environment": "example-updated"}, + ) + # Only fields in the update mask change. + update_mask = field_mask_pb2.FieldMask(paths=["description", "labels"]) + + try: + operation = client.update_workflow(workflow=workflow, update_mask=update_mask) + print("Waiting for the update to finish...") + response = operation.result(timeout=300) + except exceptions.NotFound as err: + print( + f"error: workflow '{workflow_id}' not found " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + except exceptions.GoogleAPICallError as err: + print( + f"error: failed to update workflow '{workflow_id}' " + f"in location '{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + else: + print(f"Resource ID: {response.name}") + print(f"Description: {response.description}") + print(f"Labels: {dict(response.labels)}") + print(f"Revision ID: {response.revision_id}") + + +# [END workflows_update_workflow] + + +if __name__ == "__main__": + import argparse + import os + + # Each flag defaults to its environment variable, then to an example value. + parser = argparse.ArgumentParser( + description=update_workflow.__doc__.splitlines()[0] + ) + parser.add_argument( + "--project-id", + default=os.environ.get("GOOGLE_CLOUD_PROJECT", "example-project-id"), + help="Google Cloud project ID (default: $GOOGLE_CLOUD_PROJECT)", + ) + parser.add_argument( + "--location-id", + default=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"), + help="Region of the workflow (default: $GOOGLE_CLOUD_LOCATION)", + ) + parser.add_argument( + "--workflow-id", + default=os.environ.get("WORKFLOW_ID", "myFirstWorkflow"), + help="ID of the workflow (default: $WORKFLOW_ID)", + ) + args = parser.parse_args() + + update_workflow( + project_id=args.project_id, + location_id=args.location_id, + workflow_id=args.workflow_id, + ) diff --git a/workflows/deduplicate_scheduled_runs/.python-version b/workflows/deduplicate_scheduled_runs/.python-version new file mode 100644 index 00000000000..24ee5b1be99 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/workflows/deduplicate_scheduled_runs/__init__.py b/workflows/deduplicate_scheduled_runs/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/deduplicate_scheduled_runs/main.py b/workflows/deduplicate_scheduled_runs/main.py new file mode 100644 index 00000000000..1b21b7b595f --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/main.py @@ -0,0 +1,128 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_deduplicate_scheduled_runs] +import datetime +import hashlib +import json +import os +import uuid + +import flask +import functions_framework +from google.cloud import firestore +from google.cloud.workflows import executions_v1 + +# Longer than starting an execution takes, so only a crashed attempt's claim expires. +LEASE = datetime.timedelta(minutes=5) +# Keep records past the job's retry window. A TTL policy on expire_at deletes them. +RETENTION = datetime.timedelta(days=1) +# The workflow to start, as projects/PROJECT/locations/LOCATION/workflows/WORKFLOW. +WORKFLOW_NAME = os.environ.get("WORKFLOW_NAME", "") + +db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)")) +executions_client = executions_v1.ExecutionsClient() + + +@firestore.transactional +def claim_run( + transaction: firestore.Transaction, + run_ref: firestore.DocumentReference, + attempt_id: str, + now: datetime.datetime, +) -> str: + """Claims the run unless it's done or another attempt holds the lease.""" + snapshot = run_ref.get(transaction=transaction) + if snapshot.exists: + run = snapshot.to_dict() + if run["status"] == "completed": + return "completed" + if run["lease_expires"] > now: + return "in_progress" + transaction.set( + run_ref, + { + "status": "in_progress", + "attempt_id": attempt_id, + "lease_expires": now + LEASE, + "expire_at": now + RETENTION, + }, + ) + return "claimed" + + +@firestore.transactional +def finish_run( + transaction: firestore.Transaction, + run_ref: firestore.DocumentReference, + attempt_id: str, + execution_name: str | None, +) -> None: + """Records the execution, or releases the claim, if this attempt holds it.""" + snapshot = run_ref.get(transaction=transaction) + if not snapshot.exists or snapshot.to_dict().get("attempt_id") != attempt_id: + return # A retry claimed the run after this attempt's lease expired. + if execution_name: + transaction.update( + run_ref, {"status": "completed", "execution_name": execution_name} + ) + else: + transaction.delete(run_ref) + + +@functions_framework.http +def start_scheduled_run(request: flask.Request) -> tuple[str, int]: + """Starts the workflow once per Cloud Scheduler run, even if the run is delivered twice.""" + job_name = request.headers.get("X-CloudScheduler-JobName") + schedule_time = request.headers.get("X-CloudScheduler-ScheduleTime") + if not job_name or not schedule_time: + return "error: missing Cloud Scheduler headers", 400 + + # Scheduler adds a per-job offset of a few seconds to the schedule time, but + # a forced run before the job's first scheduled run sends the time without + # it. Keying on the scheduled minute matches both. A forced run uses the + # next scheduled time, so that scheduled run is then skipped. + try: + scheduled = datetime.datetime.fromisoformat(schedule_time) + except ValueError: + return "error: invalid X-CloudScheduler-ScheduleTime header", 400 + run_minute = scheduled.astimezone(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%MZ") + run_id = hashlib.sha256(f"{job_name}/{run_minute}".encode()).hexdigest() + run_ref = db.collection("scheduled_workflow_runs").document(run_id) + attempt_id = uuid.uuid4().hex + now = datetime.datetime.now(datetime.timezone.utc) + + status = claim_run(db.transaction(), run_ref, attempt_id, now) + if status == "completed": + return "Duplicate run skipped", 200 + if status == "in_progress": + # A non-2xx response keeps retries going in case that attempt fails. + return "Run in progress", 409 + + # Pass the run's identity to the workflow, so its steps can use it as an + # idempotency key too. + argument = json.dumps({"job_name": job_name, "scheduled_minute": run_minute}) + try: + execution = executions_client.create_execution( + parent=WORKFLOW_NAME, execution=executions_v1.Execution(argument=argument) + ) + except Exception: + # Releases the claim so that a retry can start the run. + finish_run(db.transaction(), run_ref, attempt_id, None) + raise + finish_run(db.transaction(), run_ref, attempt_id, execution.name) + return f"Started execution {execution.name}", 200 + + +# [END workflows_deduplicate_scheduled_runs] diff --git a/workflows/deduplicate_scheduled_runs/noxfile_config.py b/workflows/deduplicate_scheduled_runs/noxfile_config.py new file mode 100644 index 00000000000..95c2d9ac875 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/noxfile_config.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be imported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/workflows/deduplicate_scheduled_runs/requirements-test.txt b/workflows/deduplicate_scheduled_runs/requirements-test.txt new file mode 100644 index 00000000000..f48d9c027c7 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/requirements-test.txt @@ -0,0 +1,2 @@ +pytest==9.0.3 +google-cloud-scheduler==2.21.0 diff --git a/workflows/deduplicate_scheduled_runs/requirements.txt b/workflows/deduplicate_scheduled_runs/requirements.txt new file mode 100644 index 00000000000..9740806e07f --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/requirements.txt @@ -0,0 +1,3 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 +google-cloud-workflows==1.24.0 diff --git a/workflows/deduplicate_scheduled_runs/tests/__init__.py b/workflows/deduplicate_scheduled_runs/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/deduplicate_scheduled_runs/tests/conftest.py b/workflows/deduplicate_scheduled_runs/tests/conftest.py new file mode 100644 index 00000000000..6774c4eda2a --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/tests/conftest.py @@ -0,0 +1,23 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest import mock + +# main.py creates Firestore and Workflows clients when it's imported. Without +# credentials that fails, so unit-test runs replace the client classes. The unit +# tests then swap in fakes, and the integration tests use real clients. +if os.environ.get("RUN_INTEGRATION_TESTS") != "true": + mock.patch("google.cloud.firestore.Client").start() + mock.patch("google.cloud.workflows.executions_v1.ExecutionsClient").start() diff --git a/workflows/deduplicate_scheduled_runs/tests/integration_test.py b/workflows/deduplicate_scheduled_runs/tests/integration_test.py new file mode 100644 index 00000000000..df1763573be --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/tests/integration_test.py @@ -0,0 +1,119 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test: Cloud Scheduler → this function → a workflow execution. + +Skipped unless RUN_INTEGRATION_TESTS is "true". Needs the function (main.py) +and a target workflow deployed, and: + GOOGLE_CLOUD_PROJECT, GOOGLE_CLOUD_LOCATION (default 'us-central1'), + FUNCTION_URL, SERVICE_ACCOUNT_EMAIL, WORKFLOW_ID, and FIRESTORE_DATABASE. + +A forced run sends the job's next scheduled time, so forcing the same job twice +delivers the same scheduled run twice. +""" + +import json +import os +import time +from typing import Iterator +import uuid + +from google.api_core import exceptions +from google.api_core import retry +from google.cloud import scheduler_v1 +from google.cloud.workflows import executions_v1 +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_INTEGRATION_TESTS") != "true", + reason="Set RUN_INTEGRATION_TESTS=true to run integration tests.", +) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "") +LOCATION_ID = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") +FUNCTION_URL = os.environ.get("FUNCTION_URL", "") +SERVICE_ACCOUNT_EMAIL = os.environ.get("SERVICE_ACCOUNT_EMAIL", "") +WORKFLOW_ID = os.environ.get("WORKFLOW_ID", "") + +# Cloud Scheduler rejects a change to a job while another change is in progress. +retry_aborted = retry.Retry( + predicate=retry.if_exception_type(exceptions.Aborted), timeout=120 +) + + +@pytest.fixture +def job_name() -> Iterator[str]: + """Creates a Scheduler job that calls the function, and deletes it afterward.""" + client = scheduler_v1.CloudSchedulerClient() + parent = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + name = f"{parent}/jobs/test-job-{uuid.uuid4().hex[:12]}" + client.create_job( + parent=parent, + job=scheduler_v1.Job( + name=name, + http_target=scheduler_v1.HttpTarget( + uri=FUNCTION_URL, + http_method=scheduler_v1.HttpMethod.POST, + oidc_token=scheduler_v1.OidcToken( + service_account_email=SERVICE_ACCOUNT_EMAIL, + audience=FUNCTION_URL, + ), + ), + schedule="0 9 * * 1", + time_zone="UTC", + ), + retry=retry_aborted, + ) + yield name + client.delete_job(name=name, retry=retry_aborted) + + +def executions_for(job_id: str) -> list[executions_v1.Execution]: + client = executions_v1.ExecutionsClient() + request = executions_v1.ListExecutionsRequest( + parent=f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{WORKFLOW_ID}", + view=executions_v1.ExecutionView.FULL, + ) + return [ + execution + for execution in client.list_executions(request=request) + if json.loads(execution.argument or "{}").get("job_name") == job_id + ] + + +def test_duplicate_delivery_starts_one_execution(job_name: str) -> None: + client = scheduler_v1.CloudSchedulerClient() + job_id = job_name.rsplit("/", 1)[-1] + + client.run_job(name=job_name, retry=retry_aborted) + deadline = time.monotonic() + 300 + while not executions_for(job_id): + assert time.monotonic() < deadline, "The workflow didn't start." + time.sleep(10) + + # Force the same job again: it delivers the same scheduled run. + client.run_job(name=job_name, retry=retry_aborted) + deadline = time.monotonic() + 300 + job = client.get_job(name=job_name) + first_attempt = job.last_attempt_time + while job.last_attempt_time == first_attempt: + assert time.monotonic() < deadline, "The second delivery didn't happen." + time.sleep(10) + job = client.get_job(name=job_name) + + assert job.status.code == 0, job.status.message + executions = executions_for(job_id) + assert len(executions) == 1 + argument = json.loads(executions[0].argument) + assert argument["scheduled_minute"] diff --git a/workflows/deduplicate_scheduled_runs/tests/main_test.py b/workflows/deduplicate_scheduled_runs/tests/main_test.py new file mode 100644 index 00000000000..5f7c5b06946 --- /dev/null +++ b/workflows/deduplicate_scheduled_runs/tests/main_test.py @@ -0,0 +1,201 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import json +from typing import Any +from unittest import mock + +import flask +from google.api_core import exceptions +from google.cloud.workflows import executions_v1 +import pytest + +from .. import main + +app = flask.Flask(__name__) +HEADERS = { + "X-CloudScheduler-JobName": "test-job", + "X-CloudScheduler-ScheduleTime": "2026-11-09T09:00:00.123456-08:00", +} +EXECUTION_NAME = "projects/p/locations/l/workflows/w/executions/e" + + +class FakeSnapshot: + def __init__(self, data: dict[str, Any] | None) -> None: + self._data = data + self.exists = data is not None + + def to_dict(self) -> dict[str, Any]: + return dict(self._data or {}) + + +class FakeDocument: + def __init__(self, store: dict[str, dict[str, Any]], doc_id: str) -> None: + self.store = store + self.id = doc_id + + def get(self, transaction: object = None) -> FakeSnapshot: + return FakeSnapshot(self.store.get(self.id)) + + +class FakeDatabase: + def __init__(self) -> None: + self.store: dict[str, dict[str, Any]] = {} + + def collection(self, name: str) -> mock.MagicMock: + collection = mock.MagicMock() + collection.document.side_effect = lambda doc_id: FakeDocument( + self.store, doc_id + ) + return collection + + def transaction(self) -> mock.MagicMock: + # firestore.transactional begins, commits, and retries the transaction. + transaction = mock.MagicMock(_max_attempts=1, _id=b"test", _read_only=False) + transaction.set.side_effect = lambda ref, data: self.store.__setitem__( + ref.id, dict(data) + ) + transaction.update.side_effect = lambda ref, data: self.store[ref.id].update( + data + ) + transaction.delete.side_effect = lambda ref: self.store.pop(ref.id, None) + return transaction + + +@pytest.fixture +def db(monkeypatch: pytest.MonkeyPatch) -> FakeDatabase: + database = FakeDatabase() + monkeypatch.setattr(main, "db", database) + return database + + +@pytest.fixture +def executions(monkeypatch: pytest.MonkeyPatch) -> mock.MagicMock: + client = mock.MagicMock() + client.create_execution.return_value = executions_v1.Execution(name=EXECUTION_NAME) + monkeypatch.setattr(main, "executions_client", client) + monkeypatch.setattr(main, "WORKFLOW_NAME", "projects/p/locations/l/workflows/w") + return client + + +def call(headers: dict[str, str] | None = None) -> tuple[str, int]: + with app.test_request_context("/", method="POST", headers=headers): + return main.start_scheduled_run(flask.request) + + +def only_record(db: FakeDatabase) -> dict[str, Any]: + assert len(db.store) == 1 + return next(iter(db.store.values())) + + +def test_first_delivery_starts_the_workflow( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + body, status = call(HEADERS) + + assert status == 200 + kwargs = executions.create_execution.call_args.kwargs + assert kwargs["parent"] == "projects/p/locations/l/workflows/w" + assert json.loads(kwargs["execution"].argument) == { + "job_name": "test-job", + "scheduled_minute": "2026-11-09T17:00Z", + } + record = only_record(db) + assert record["status"] == "completed" + assert record["execution_name"] == EXECUTION_NAME + + +def test_duplicate_delivery_is_skipped( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + call(HEADERS) + + body, status = call(HEADERS) + + assert (body, status) == ("Duplicate run skipped", 200) + executions.create_execution.assert_called_once() + + +def test_schedule_time_with_and_without_offset_is_one_run( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + # A forced run before the job's first scheduled run sends the time without + # Scheduler's per-job offset; later deliveries include it. + call( + dict(HEADERS, **{"X-CloudScheduler-ScheduleTime": "2026-11-09T09:00:00-08:00"}) + ) + + body, status = call(HEADERS) + + assert (body, status) == ("Duplicate run skipped", 200) + executions.create_execution.assert_called_once() + + +def test_invalid_schedule_time_is_rejected( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + body, status = call(dict(HEADERS, **{"X-CloudScheduler-ScheduleTime": "soon"})) + + assert status == 400 + executions.create_execution.assert_not_called() + + +def test_delivery_with_active_claim_returns_409( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + call(HEADERS) + future = datetime.datetime.now(datetime.timezone.utc) + main.LEASE + only_record(db).update(status="in_progress", lease_expires=future) + + _, status = call(HEADERS) + + assert status == 409 + executions.create_execution.assert_called_once() + + +def test_failed_start_releases_the_claim( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + executions.create_execution.side_effect = exceptions.ServiceUnavailable("down") + + with pytest.raises(exceptions.ServiceUnavailable): + call(HEADERS) + + assert db.store == {} + + +def test_stale_attempt_does_not_complete_newer_claim( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + def slow_start(**kwargs: object) -> executions_v1.Execution: + # While this attempt runs, its lease expires and a retry claims the run. + only_record(db).update(attempt_id="newer-attempt") + return executions_v1.Execution(name=EXECUTION_NAME) + + executions.create_execution.side_effect = slow_start + + call(HEADERS) + + assert only_record(db)["status"] == "in_progress" + + +def test_missing_headers_are_rejected( + db: FakeDatabase, executions: mock.MagicMock +) -> None: + body, status = call() + + assert status == 400 + assert body.startswith("error:") + executions.create_execution.assert_not_called() diff --git a/workflows/idempotent_payment/.python-version b/workflows/idempotent_payment/.python-version new file mode 100644 index 00000000000..24ee5b1be99 --- /dev/null +++ b/workflows/idempotent_payment/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/workflows/idempotent_payment/__init__.py b/workflows/idempotent_payment/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/idempotent_payment/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/idempotent_payment/main.py b/workflows/idempotent_payment/main.py new file mode 100644 index 00000000000..29fc5c9b498 --- /dev/null +++ b/workflows/idempotent_payment/main.py @@ -0,0 +1,128 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_idempotent_process_payment] +import datetime +import os +import uuid + +import flask +import functions_framework +from google.auth.transport.requests import Request +from google.cloud import firestore +from google.oauth2 import id_token +import requests + +# Longer than a payment takes, so only a crashed attempt's claim expires. +LEASE = datetime.timedelta(minutes=2) +PAYMENT_SERVICE_URL = os.environ.get("PAYMENT_SERVICE_URL", "") +# Later states, such as refunded, are final: retrying the order doesn't charge again. +CHARGEABLE_STATES = ("created", "payment_failed", "payment_in_progress") + +db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)")) + + +@firestore.transactional +def claim_payment( + transaction: firestore.Transaction, + order_ref: firestore.DocumentReference, + attempt_id: str, + now: datetime.datetime, +) -> tuple[str, int]: + """Claims the order's payment unless it's paid or another attempt holds it. + + Returns the claim status and the order's amount. + """ + snapshot = order_ref.get(transaction=transaction) + if not snapshot.exists: + return "missing", 0 + order = snapshot.to_dict() + if order["order_state"] == "paid": + return "paid", order["amount_cents"] + if order["order_state"] == "payment_in_progress" and order["lease_expires"] > now: + return "in_progress", order["amount_cents"] + if order["order_state"] not in CHARGEABLE_STATES: + return "final", order["amount_cents"] + transaction.update( + order_ref, + { + "order_state": "payment_in_progress", + "attempt_id": attempt_id, + "lease_expires": now + LEASE, + }, + ) + return "claimed", order["amount_cents"] + + +@firestore.transactional +def finish_payment( + transaction: firestore.Transaction, + order_ref: firestore.DocumentReference, + attempt_id: str, + order_state: str, +) -> None: + """Records the result, if this attempt still holds the claim.""" + snapshot = order_ref.get(transaction=transaction) + if snapshot.exists and snapshot.to_dict().get("attempt_id") == attempt_id: + transaction.update(order_ref, {"order_state": order_state}) + + +def charge(order_id: str, amount_cents: int) -> None: + """Charges the order through the payment service, at most once per order ID.""" + token = id_token.fetch_id_token(Request(), PAYMENT_SERVICE_URL) + response = requests.post( + PAYMENT_SERVICE_URL, + json={"order_id": order_id, "amount_cents": amount_cents}, + # The payment service charges each idempotency key only once, so a + # retry after an ambiguous failure can't charge the customer twice. + headers={"Authorization": f"Bearer {token}", "Idempotency-Key": order_id}, + timeout=30, + ) + response.raise_for_status() + + +@functions_framework.http +def process_payment(request: flask.Request) -> tuple[dict, int]: + """Processes an order's payment once, even when the workflow retries the call.""" + order_id = (request.get_json(silent=True) or {}).get("order_id") + if not order_id: + return {"error": "missing order_id"}, 400 + + order_ref = db.collection("orders").document(str(order_id)) + attempt_id = uuid.uuid4().hex + now = datetime.datetime.now(datetime.timezone.utc) + + status, amount_cents = claim_payment(db.transaction(), order_ref, attempt_id, now) + if status == "missing": + return {"error": f"order {order_id} not found"}, 404 + if status == "paid": + return {"order_id": order_id, "status": "paid", "already_processed": True}, 200 + if status == "in_progress": + # The workflow's retry policy retries 503 responses after a backoff. + return {"error": "payment in progress"}, 503 + if status == "final": + return {"error": f"order {order_id} can't be charged again"}, 409 + + # Charge outside the transaction: Firestore can rerun a transaction, but it + # can't undo a call to another service. + try: + charge(order_id, amount_cents) + except requests.RequestException as err: + finish_payment(db.transaction(), order_ref, attempt_id, "payment_failed") + return {"error": f"payment failed: {err}"}, 503 + finish_payment(db.transaction(), order_ref, attempt_id, "paid") + return {"order_id": order_id, "status": "paid", "already_processed": False}, 200 + + +# [END workflows_idempotent_process_payment] diff --git a/workflows/idempotent_payment/noxfile_config.py b/workflows/idempotent_payment/noxfile_config.py new file mode 100644 index 00000000000..95c2d9ac875 --- /dev/null +++ b/workflows/idempotent_payment/noxfile_config.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be imported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/workflows/idempotent_payment/requirements-test.txt b/workflows/idempotent_payment/requirements-test.txt new file mode 100644 index 00000000000..b3b2b3ca4cc --- /dev/null +++ b/workflows/idempotent_payment/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3 diff --git a/workflows/idempotent_payment/requirements.txt b/workflows/idempotent_payment/requirements.txt new file mode 100644 index 00000000000..deda6b36e7c --- /dev/null +++ b/workflows/idempotent_payment/requirements.txt @@ -0,0 +1,3 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 +requests==2.34.2 diff --git a/workflows/idempotent_payment/tests/__init__.py b/workflows/idempotent_payment/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/idempotent_payment/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/idempotent_payment/tests/conftest.py b/workflows/idempotent_payment/tests/conftest.py new file mode 100644 index 00000000000..b68ff8d1eb5 --- /dev/null +++ b/workflows/idempotent_payment/tests/conftest.py @@ -0,0 +1,22 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest import mock + +# main.py creates a Firestore client when it's imported. Without credentials +# that fails, so unit-test runs replace the client class. The unit tests then +# swap in a fake database, and the integration tests use a real one. +if os.environ.get("RUN_INTEGRATION_TESTS") != "true": + mock.patch("google.cloud.firestore.Client").start() diff --git a/workflows/idempotent_payment/tests/integration_test.py b/workflows/idempotent_payment/tests/integration_test.py new file mode 100644 index 00000000000..bc5e1c1352c --- /dev/null +++ b/workflows/idempotent_payment/tests/integration_test.py @@ -0,0 +1,182 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fault-injection tests that run the order workflow end to end. + +These tests are skipped unless RUN_INTEGRATION_TESTS is "true". They need the +workflow (workflow.yaml), the payment function (main.py), and the stand-in +payment service (tests/payment_service) deployed, and: + GOOGLE_CLOUD_PROJECT, GOOGLE_CLOUD_LOCATION (default 'us-central1'), + WORKFLOW_ID (default 'idempotent-order-workflow'), and FIRESTORE_DATABASE. + +Each test checks how many times the stand-in payment service charged. +""" + +import concurrent.futures +import datetime +import json +import os +import time +from typing import Any, Iterator +import uuid + +from google.cloud import firestore +from google.cloud.workflows import executions_v1 +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_INTEGRATION_TESTS") != "true", + reason="Set RUN_INTEGRATION_TESTS=true to run integration tests.", +) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "") +LOCATION_ID = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") +WORKFLOW_ID = os.environ.get("WORKFLOW_ID", "idempotent-order-workflow") +DATABASE = os.environ.get("FIRESTORE_DATABASE", "(default)") +AMOUNT_CENTS = 1999 + + +@pytest.fixture(scope="module") +def db() -> firestore.Client: + return firestore.Client(database=DATABASE) + + +@pytest.fixture +def order_id(db: firestore.Client) -> Iterator[str]: + """Provides a unique order ID and deletes its test records afterward.""" + order_id = f"test-order-{uuid.uuid4().hex[:12]}" + yield order_id + for collection in ("orders", "test_payment_charges", "test_payment_faults"): + db.collection(collection).document(order_id).delete() + + +def run_workflow(order_id: str) -> executions_v1.Execution: + """Runs the workflow for an order and waits for the execution to finish.""" + client = executions_v1.ExecutionsClient() + parent = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{WORKFLOW_ID}" + argument = json.dumps({"order_id": order_id, "amount_cents": AMOUNT_CENTS}) + execution = client.create_execution( + parent=parent, execution=executions_v1.Execution(argument=argument) + ) + deadline = time.monotonic() + 600 + while execution.state == executions_v1.Execution.State.ACTIVE: + assert time.monotonic() < deadline, "The execution didn't finish." + time.sleep(3) + execution = client.get_execution(name=execution.name) + assert execution.state == executions_v1.Execution.State.SUCCEEDED, execution.error + return execution + + +def charge_record(db: firestore.Client, order_id: str) -> dict[str, Any]: + snapshot = db.collection("test_payment_charges").document(order_id).get() + assert snapshot.exists, "The payment service wasn't called." + return snapshot.to_dict() + + +def order_state(db: firestore.Client, order_id: str) -> str: + return db.collection("orders").document(order_id).get().to_dict()["order_state"] + + +def test_duplicate_execution_charges_once(db: firestore.Client, order_id: str) -> None: + first = json.loads(run_workflow(order_id).result) + second = json.loads(run_workflow(order_id).result) + + assert first["already_processed"] is False + assert second["already_processed"] is True + assert charge_record(db, order_id)["charges"] == 1 + assert order_state(db, order_id) == "paid" + + +def test_client_retry_resumes_after_partial_failure( + db: firestore.Client, order_id: str +) -> None: + # A previous execution created the order and then failed before payment. + db.collection("orders").document(order_id).set( + {"order_state": "created", "amount_cents": AMOUNT_CENTS} + ) + + run_workflow(order_id) + + assert order_state(db, order_id) == "paid" + assert charge_record(db, order_id)["charges"] == 1 + + +def test_ambiguous_payment_failure_charges_once( + db: firestore.Client, order_id: str +) -> None: + # The payment service charges, then responds with a gateway timeout. + db.collection("test_payment_faults").document(order_id).set( + {"mode": "fail_after_charge"} + ) + + run_workflow(order_id) + + record = charge_record(db, order_id) + assert record["charges"] == 1 + assert record["requests"] >= 2 + assert order_state(db, order_id) == "paid" + + +def test_timeouts_at_two_tiers_charge_once(db: firestore.Client, order_id: str) -> None: + # The payment service charges, then responds after both the payment + # function and the workflow stop waiting. + db.collection("test_payment_faults").document(order_id).set( + {"mode": "slow_after_charge", "seconds": 45} + ) + + run_workflow(order_id) + + assert charge_record(db, order_id)["charges"] == 1 + assert order_state(db, order_id) == "paid" + + +def test_crashed_attempt_is_taken_over(db: firestore.Client, order_id: str) -> None: + # An earlier attempt claimed the payment and crashed; its claim expired. + expired = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + minutes=5 + ) + db.collection("orders").document(order_id).set( + { + "order_state": "payment_in_progress", + "amount_cents": AMOUNT_CENTS, + "attempt_id": "crashed-attempt", + "lease_expires": expired, + } + ) + + run_workflow(order_id) + + assert order_state(db, order_id) == "paid" + assert charge_record(db, order_id)["charges"] == 1 + + +def test_concurrent_executions_charge_once(db: firestore.Client, order_id: str) -> None: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(run_workflow, [order_id, order_id])) + + outcomes = sorted(json.loads(r.result)["already_processed"] for r in results) + assert outcomes == [False, True] + assert charge_record(db, order_id)["charges"] == 1 + + +def test_order_id_reaches_the_payment_service( + db: firestore.Client, order_id: str +) -> None: + run_workflow(order_id) + + # The stand-in stores charges by Idempotency-Key, and records the order ID + # from the request body. + record = charge_record(db, order_id) + assert record["order_id"] == order_id + assert record["amount_cents"] == AMOUNT_CENTS diff --git a/workflows/idempotent_payment/tests/main_test.py b/workflows/idempotent_payment/tests/main_test.py new file mode 100644 index 00000000000..234317c764e --- /dev/null +++ b/workflows/idempotent_payment/tests/main_test.py @@ -0,0 +1,224 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +from typing import Any +from unittest import mock + +import flask +import pytest +import requests + +from .. import main + +app = flask.Flask(__name__) +ORDER_ID = "order-1001" +NOW = datetime.datetime.now(datetime.timezone.utc) + + +class FakeSnapshot: + def __init__(self, data: dict[str, Any] | None) -> None: + self._data = data + self.exists = data is not None + + def to_dict(self) -> dict[str, Any]: + return dict(self._data or {}) + + +class FakeDocument: + def __init__(self, store: dict[str, dict[str, Any]], doc_id: str) -> None: + self.store = store + self.id = doc_id + + def get(self, transaction: object = None) -> FakeSnapshot: + return FakeSnapshot(self.store.get(self.id)) + + +class FakeDatabase: + """An in-memory stand-in for the orders collection.""" + + def __init__(self) -> None: + self.store: dict[str, dict[str, Any]] = {} + + def collection(self, name: str) -> mock.MagicMock: + collection = mock.MagicMock() + collection.document.side_effect = lambda doc_id: FakeDocument( + self.store, doc_id + ) + return collection + + def transaction(self) -> mock.MagicMock: + # firestore.transactional begins, commits, and retries the transaction. + transaction = mock.MagicMock(_max_attempts=1, _id=b"test", _read_only=False) + transaction.update.side_effect = lambda ref, data: self.store[ref.id].update( + data + ) + return transaction + + +@pytest.fixture +def db(monkeypatch: pytest.MonkeyPatch) -> FakeDatabase: + database = FakeDatabase() + monkeypatch.setattr(main, "db", database) + return database + + +@pytest.fixture +def charge(monkeypatch: pytest.MonkeyPatch) -> mock.MagicMock: + charge_mock = mock.MagicMock() + monkeypatch.setattr(main, "charge", charge_mock) + return charge_mock + + +def create_order(db: FakeDatabase, **fields: object) -> dict[str, Any]: + db.store[ORDER_ID] = {"order_state": "created", "amount_cents": 1999, **fields} + return db.store[ORDER_ID] + + +def call(body: dict[str, Any] | None = None) -> tuple[dict, int]: + body = {"order_id": ORDER_ID} if body is None else body + with app.test_request_context("/", method="POST", json=body): + return main.process_payment(flask.request) + + +def test_new_order_is_charged_once(db: FakeDatabase, charge: mock.MagicMock) -> None: + order = create_order(db) + + body, status = call() + + assert status == 200 + assert body["already_processed"] is False + charge.assert_called_once_with(ORDER_ID, 1999) + assert order["order_state"] == "paid" + + +def test_paid_order_is_not_charged_again( + db: FakeDatabase, charge: mock.MagicMock +) -> None: + create_order(db, order_state="paid") + + body, status = call() + + assert (status, body["already_processed"]) == (200, True) + charge.assert_not_called() + + +def test_refunded_order_is_not_charged_again( + db: FakeDatabase, charge: mock.MagicMock +) -> None: + order = create_order(db, order_state="refunded") + + _, status = call() + + assert status == 409 + charge.assert_not_called() + assert order["order_state"] == "refunded" + + +def test_order_with_active_claim_returns_503( + db: FakeDatabase, charge: mock.MagicMock +) -> None: + create_order( + db, + order_state="payment_in_progress", + attempt_id="other-attempt", + lease_expires=NOW + main.LEASE, + ) + + _, status = call() + + assert status == 503 + charge.assert_not_called() + + +@pytest.mark.parametrize( + "fields", + [ + { + "order_state": "payment_in_progress", + "attempt_id": "crashed-attempt", + "lease_expires": NOW - datetime.timedelta(seconds=1), + }, + {"order_state": "payment_failed"}, + ], + ids=["expired_claim", "failed_payment"], +) +def test_unfinished_order_is_charged( + db: FakeDatabase, charge: mock.MagicMock, fields: dict[str, Any] +) -> None: + order = create_order(db, **fields) + + _, status = call() + + assert status == 200 + charge.assert_called_once() + assert order["order_state"] == "paid" + + +def test_ambiguous_payment_failure_is_retryable( + db: FakeDatabase, charge: mock.MagicMock +) -> None: + order = create_order(db) + charge.side_effect = requests.Timeout("timed out") + + _, status = call() + + assert status == 503 + assert order["order_state"] == "payment_failed" + + +def test_stale_attempt_does_not_overwrite_newer_claim( + db: FakeDatabase, charge: mock.MagicMock +) -> None: + order = create_order(db) + # While this attempt charges, its lease expires and a retry claims the order. + charge.side_effect = lambda *args: order.update(attempt_id="newer-attempt") + + call() + + assert order["order_state"] == "payment_in_progress" + assert order["attempt_id"] == "newer-attempt" + + +@pytest.mark.parametrize( + "body, expected_status", + [({}, 400), ({"order_id": "missing-order"}, 404)], +) +def test_bad_requests( + db: FakeDatabase, + charge: mock.MagicMock, + body: dict[str, Any], + expected_status: int, +) -> None: + _, status = call(body) + + assert status == expected_status + charge.assert_not_called() + + +def test_charge_sends_order_id_as_idempotency_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(main, "PAYMENT_SERVICE_URL", "https://payments.example.com") + monkeypatch.setattr(main.id_token, "fetch_id_token", lambda *args: "token") + post = mock.MagicMock() + monkeypatch.setattr(main.requests, "post", post) + + main.charge(ORDER_ID, 1999) + + kwargs = post.call_args.kwargs + assert post.call_args.args == ("https://payments.example.com",) + assert kwargs["headers"]["Idempotency-Key"] == ORDER_ID + assert kwargs["headers"]["Authorization"] == "Bearer token" + assert kwargs["json"] == {"order_id": ORDER_ID, "amount_cents": 1999} diff --git a/workflows/idempotent_payment/tests/payment_service/main.py b/workflows/idempotent_payment/tests/payment_service/main.py new file mode 100644 index 00000000000..1f258ef0a37 --- /dev/null +++ b/workflows/idempotent_payment/tests/payment_service/main.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A stand-in payment service for testing, not a published sample. + +It charges each Idempotency-Key once and records every request in Firestore, +so tests can count charges. A fault document for a key makes the next request +charge and then respond ambiguously (HTTP 504), or charge and respond slowly. +""" + +import os +import time +import uuid + +import flask +import functions_framework +from google.cloud import firestore + +db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)")) + + +@firestore.transactional +def record_charge( + transaction: firestore.Transaction, + charge_ref: firestore.DocumentReference, + order_id: str, + amount_cents: int, +) -> bool: + """Records one charge per key. Returns True if this request charged.""" + snapshot = charge_ref.get(transaction=transaction) + if snapshot.exists: + transaction.update(charge_ref, {"requests": firestore.Increment(1)}) + return False + transaction.set( + charge_ref, + { + "charge_id": uuid.uuid4().hex, + "order_id": order_id, + "amount_cents": amount_cents, + "charges": 1, + "requests": 1, + }, + ) + return True + + +@functions_framework.http +def charge(request: flask.Request) -> tuple[dict, int]: + key = request.headers.get("Idempotency-Key") + body = request.get_json(silent=True) or {} + if not key: + return {"error": "missing Idempotency-Key"}, 400 + + charge_ref = db.collection("test_payment_charges").document(key) + record_charge( + db.transaction(), charge_ref, body.get("order_id"), body.get("amount_cents") + ) + + fault_ref = db.collection("test_payment_faults").document(key) + fault = fault_ref.get() + if fault.exists: + fault_ref.delete() # One-shot: the next request behaves normally. + mode = fault.to_dict().get("mode") + if mode == "fail_after_charge": + return {"error": "gateway timeout"}, 504 + if mode == "slow_after_charge": + time.sleep(fault.to_dict().get("seconds", 40)) + return {"charge_id": charge_ref.get().to_dict()["charge_id"]}, 200 diff --git a/workflows/idempotent_payment/tests/payment_service/requirements.txt b/workflows/idempotent_payment/tests/payment_service/requirements.txt new file mode 100644 index 00000000000..6d4f99e4b72 --- /dev/null +++ b/workflows/idempotent_payment/tests/payment_service/requirements.txt @@ -0,0 +1,2 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 diff --git a/workflows/idempotent_payment/workflow.yaml b/workflows/idempotent_payment/workflow.yaml new file mode 100644 index 00000000000..3cd32053b44 --- /dev/null +++ b/workflows/idempotent_payment/workflow.yaml @@ -0,0 +1,83 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_idempotent_order_workflow] +# Processes an order once per client-supplied order_id, even when the client +# retries the execution or the workflow retries the payment call. +# Deploy with the PAYMENT_FUNCTION_URL environment variable, and optionally +# FIRESTORE_DATABASE. +main: + params: [args] + steps: + - init: + assign: + - order_id: ${string(args.order_id)} + - database: ${default(sys.get_env("FIRESTORE_DATABASE"), "(default)")} + - documents: ${"projects/" + sys.get_env("GOOGLE_CLOUD_PROJECT_ID") + "/databases/" + database + "/documents"} + - create_order: + try: + # Creating a document with the order ID fails with 409 if it exists, + # so concurrent executions can't both create the order. + call: googleapis.firestore.v1.projects.databases.documents.createDocument + args: + collectionId: orders + parent: ${documents} + documentId: ${order_id} + body: + fields: + order_state: + stringValue: created + amount_cents: + integerValue: ${string(args.amount_cents)} + except: + as: e + steps: + - check_conflict: + switch: + - condition: ${map.get(e, "code") == 409} + next: read_order + - unhandled_error: + raise: ${e} + next: call_payment + - read_order: + # A retried request finds the existing order and resumes from its state, + # instead of exiting and leaving an unpaid order behind. + call: googleapis.firestore.v1.projects.databases.documents.get + args: + name: ${documents + "/orders/" + order_id} + result: existing_order + - check_state: + switch: + - condition: ${existing_order.fields.order_state.stringValue == "paid"} + return: + order_id: ${order_id} + status: paid + already_processed: true + - call_payment: + try: + call: http.post + args: + url: ${sys.get_env("PAYMENT_FUNCTION_URL")} + auth: + type: OIDC + body: + order_id: ${order_id} + timeout: 30 + result: payment + # The payment function is idempotent, so retrying timeouts and + # 429, 502, 503, and 504 responses can't charge twice. + retry: ${http.default_retry} + - finish: + return: ${payment.body} +# [END workflows_idempotent_order_workflow] diff --git a/workflows/order_saga/.python-version b/workflows/order_saga/.python-version new file mode 100644 index 00000000000..24ee5b1be99 --- /dev/null +++ b/workflows/order_saga/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/workflows/order_saga/__init__.py b/workflows/order_saga/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/order_saga/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/order_saga/main.py b/workflows/order_saga/main.py new file mode 100644 index 00000000000..227b4155224 --- /dev/null +++ b/workflows/order_saga/main.py @@ -0,0 +1,135 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_order_saga_refund] +import datetime +import os +import uuid + +import flask +import functions_framework +from google.auth.transport.requests import Request +from google.cloud import firestore +from google.oauth2 import id_token +import requests + +# Longer than a refund takes, so only a crashed attempt's claim expires. +LEASE = datetime.timedelta(minutes=2) +PAYMENT_SERVICE_URL = os.environ.get("PAYMENT_SERVICE_URL", "") +REFUNDABLE_STATES = ("paid", "refund_failed") + +db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)")) + + +@firestore.transactional +def claim_refund( + transaction: firestore.Transaction, + order_ref: firestore.DocumentReference, + attempt_id: str, + now: datetime.datetime, +) -> tuple[str, int]: + """Claims the order's refund unless it's refunded or another attempt holds it. + + Returns the claim status and the order's amount. + """ + snapshot = order_ref.get(transaction=transaction) + if not snapshot.exists: + return "missing", 0 + order = snapshot.to_dict() + state = order["order_state"] + if state == "refunded": + return "refunded", order["amount_cents"] + if state == "refund_in_progress" and order["lease_expires"] > now: + return "in_progress", order["amount_cents"] + if state not in REFUNDABLE_STATES and state != "refund_in_progress": + return "not_paid", order["amount_cents"] + transaction.update( + order_ref, + { + "order_state": "refund_in_progress", + "attempt_id": attempt_id, + "lease_expires": now + LEASE, + }, + ) + return "claimed", order["amount_cents"] + + +@firestore.transactional +def finish_refund( + transaction: firestore.Transaction, + order_ref: firestore.DocumentReference, + attempt_id: str, + order_state: str, +) -> None: + """Records the result, if this attempt still holds the claim.""" + snapshot = order_ref.get(transaction=transaction) + if snapshot.exists and snapshot.to_dict().get("attempt_id") == attempt_id: + transaction.update(order_ref, {"order_state": order_state}) + + +def refund(order_id: str, amount_cents: int) -> None: + """Refunds the order through the payment service, at most once per order ID.""" + token = id_token.fetch_id_token(Request(), PAYMENT_SERVICE_URL) + response = requests.post( + PAYMENT_SERVICE_URL, + json={"order_id": order_id, "amount_cents": -amount_cents}, + # A key that's different from the charge's key, so the payment service + # treats the refund as its own operation, also done only once. + headers={ + "Authorization": f"Bearer {token}", + "Idempotency-Key": f"refund-{order_id}", + }, + timeout=30, + ) + response.raise_for_status() + + +@functions_framework.http +def refund_order(request: flask.Request) -> tuple[dict, int]: + """Refunds a paid order once, even when the workflow retries the call.""" + order_id = (request.get_json(silent=True) or {}).get("order_id") + if not order_id: + return {"error": "missing order_id"}, 400 + + order_ref = db.collection("orders").document(str(order_id)) + attempt_id = uuid.uuid4().hex + now = datetime.datetime.now(datetime.timezone.utc) + + status, amount_cents = claim_refund(db.transaction(), order_ref, attempt_id, now) + if status == "missing": + return {"error": f"order {order_id} not found"}, 404 + if status == "not_paid": + return {"error": f"order {order_id} wasn't paid"}, 409 + if status == "refunded": + return { + "order_id": order_id, + "status": "refunded", + "already_processed": True, + }, 200 + if status == "in_progress": + # The workflow's retry policy retries 503 responses after a backoff. + return {"error": "refund in progress"}, 503 + + # Refund outside the transaction: Firestore can rerun a transaction, but it + # can't undo a call to another service. + try: + refund(order_id, amount_cents) + except requests.RequestException as err: + finish_refund(db.transaction(), order_ref, attempt_id, "refund_failed") + return {"error": f"refund failed: {err}"}, 503 + finish_refund(db.transaction(), order_ref, attempt_id, "refunded") + return {"order_id": order_id, "status": "refunded", "already_processed": False}, 200 + + +# [END workflows_order_saga_refund] diff --git a/workflows/order_saga/noxfile_config.py b/workflows/order_saga/noxfile_config.py new file mode 100644 index 00000000000..95c2d9ac875 --- /dev/null +++ b/workflows/order_saga/noxfile_config.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be imported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/workflows/order_saga/requirements-test.txt b/workflows/order_saga/requirements-test.txt new file mode 100644 index 00000000000..b3b2b3ca4cc --- /dev/null +++ b/workflows/order_saga/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3 diff --git a/workflows/order_saga/requirements.txt b/workflows/order_saga/requirements.txt new file mode 100644 index 00000000000..deda6b36e7c --- /dev/null +++ b/workflows/order_saga/requirements.txt @@ -0,0 +1,3 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 +requests==2.34.2 diff --git a/workflows/order_saga/tests/__init__.py b/workflows/order_saga/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/workflows/order_saga/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/workflows/order_saga/tests/conftest.py b/workflows/order_saga/tests/conftest.py new file mode 100644 index 00000000000..b68ff8d1eb5 --- /dev/null +++ b/workflows/order_saga/tests/conftest.py @@ -0,0 +1,22 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest import mock + +# main.py creates a Firestore client when it's imported. Without credentials +# that fails, so unit-test runs replace the client class. The unit tests then +# swap in a fake database, and the integration tests use a real one. +if os.environ.get("RUN_INTEGRATION_TESTS") != "true": + mock.patch("google.cloud.firestore.Client").start() diff --git a/workflows/order_saga/tests/fulfillment_service/main.py b/workflows/order_saga/tests/fulfillment_service/main.py new file mode 100644 index 00000000000..41b81aa73fb --- /dev/null +++ b/workflows/order_saga/tests/fulfillment_service/main.py @@ -0,0 +1,37 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A stand-in fulfillment service for testing, not a published sample. + +It fulfills every order unless a fault document for the order ID says to fail, +in which case it responds with HTTP 422, a permanent failure. +""" + +import os + +import flask +import functions_framework +from google.cloud import firestore + +db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)")) + + +@functions_framework.http +def fulfill(request: flask.Request) -> tuple[dict, int]: + order_id = str((request.get_json(silent=True) or {}).get("order_id", "")) + if not order_id: + return {"error": "missing order_id"}, 400 + if db.collection("test_fulfillment_faults").document(order_id).get().exists: + return {"error": "out of stock"}, 422 + return {"order_id": order_id, "status": "fulfilled"}, 200 diff --git a/workflows/order_saga/tests/fulfillment_service/requirements.txt b/workflows/order_saga/tests/fulfillment_service/requirements.txt new file mode 100644 index 00000000000..6d4f99e4b72 --- /dev/null +++ b/workflows/order_saga/tests/fulfillment_service/requirements.txt @@ -0,0 +1,2 @@ +functions-framework==3.10.2 +google-cloud-firestore==2.34.1 diff --git a/workflows/order_saga/tests/integration_test.py b/workflows/order_saga/tests/integration_test.py new file mode 100644 index 00000000000..e39ee7955dd --- /dev/null +++ b/workflows/order_saga/tests/integration_test.py @@ -0,0 +1,139 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests that run the order saga end to end. + +Skipped unless RUN_INTEGRATION_TESTS is "true". They need the saga workflow +(workflow.yaml), the refund function (main.py), the payment function from +../idempotent_payment, and the stand-in payment and fulfillment services +deployed, and: + GOOGLE_CLOUD_PROJECT, GOOGLE_CLOUD_LOCATION (default 'us-central1'), + WORKFLOW_ID (default 'order-saga-workflow'), and FIRESTORE_DATABASE. +""" + +import json +import os +import time +from typing import Any, Iterator +import uuid + +from google.cloud import firestore +from google.cloud.workflows import executions_v1 +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_INTEGRATION_TESTS") != "true", + reason="Set RUN_INTEGRATION_TESTS=true to run integration tests.", +) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "") +LOCATION_ID = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") +WORKFLOW_ID = os.environ.get("WORKFLOW_ID", "order-saga-workflow") +DATABASE = os.environ.get("FIRESTORE_DATABASE", "(default)") +AMOUNT_CENTS = 4999 + + +@pytest.fixture(scope="module") +def db() -> firestore.Client: + return firestore.Client(database=DATABASE) + + +@pytest.fixture +def order_id(db: firestore.Client) -> Iterator[str]: + """Provides a unique order ID and deletes its test records afterward.""" + order_id = f"test-order-{uuid.uuid4().hex[:12]}" + yield order_id + db.collection("orders").document(order_id).delete() + db.collection("test_fulfillment_faults").document(order_id).delete() + for key in (order_id, f"refund-{order_id}"): + db.collection("test_payment_charges").document(key).delete() + db.collection("test_payment_faults").document(key).delete() + + +def run_saga(order_id: str) -> executions_v1.Execution: + """Runs the saga for an order and waits for the execution to finish.""" + client = executions_v1.ExecutionsClient() + parent = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/workflows/{WORKFLOW_ID}" + argument = json.dumps({"order_id": order_id, "amount_cents": AMOUNT_CENTS}) + execution = client.create_execution( + parent=parent, execution=executions_v1.Execution(argument=argument) + ) + deadline = time.monotonic() + 600 + while execution.state == executions_v1.Execution.State.ACTIVE: + assert time.monotonic() < deadline, "The execution didn't finish." + time.sleep(3) + execution = client.get_execution(name=execution.name) + return execution + + +def ledger(db: firestore.Client, key: str) -> dict[str, Any] | None: + """Returns the stand-in payment service's record for a key, if any.""" + snapshot = db.collection("test_payment_charges").document(key).get() + return snapshot.to_dict() if snapshot.exists else None + + +def order_state(db: firestore.Client, order_id: str) -> str: + return db.collection("orders").document(order_id).get().to_dict()["order_state"] + + +def test_fulfilled_order_is_not_refunded(db: firestore.Client, order_id: str) -> None: + execution = run_saga(order_id) + + assert json.loads(execution.result)["status"] == "fulfilled" + assert ledger(db, order_id)["charges"] == 1 + assert ledger(db, f"refund-{order_id}") is None + + +def test_failed_fulfillment_refunds_once(db: firestore.Client, order_id: str) -> None: + db.collection("test_fulfillment_faults").document(order_id).set({"mode": "fail"}) + + execution = run_saga(order_id) + + assert json.loads(execution.result)["status"] == "refunded" + assert ledger(db, order_id)["charges"] == 1 + refund = ledger(db, f"refund-{order_id}") + assert refund["charges"] == 1 + assert refund["amount_cents"] == -AMOUNT_CENTS + assert order_state(db, order_id) == "refunded" + + +def test_retried_compensation_refunds_once(db: firestore.Client, order_id: str) -> None: + db.collection("test_fulfillment_faults").document(order_id).set({"mode": "fail"}) + # The payment service refunds, then responds with a gateway timeout. + db.collection("test_payment_faults").document(f"refund-{order_id}").set( + {"mode": "fail_after_charge"} + ) + + execution = run_saga(order_id) + + assert json.loads(execution.result)["status"] == "refunded" + refund = ledger(db, f"refund-{order_id}") + assert refund["charges"] == 1 + assert refund["requests"] >= 2 + + +def test_retrying_a_refunded_order_moves_no_money( + db: firestore.Client, order_id: str +) -> None: + db.collection("test_fulfillment_faults").document(order_id).set({"mode": "fail"}) + run_saga(order_id) + + # The client retries the same order after it was refunded. The payment + # function refuses to charge an order in a final state. + execution = run_saga(order_id) + + assert execution.state == executions_v1.Execution.State.FAILED + assert ledger(db, order_id)["charges"] == 1 + assert ledger(db, f"refund-{order_id}")["charges"] == 1 + assert order_state(db, order_id) == "refunded" diff --git a/workflows/order_saga/tests/main_test.py b/workflows/order_saga/tests/main_test.py new file mode 100644 index 00000000000..90ff646f7e8 --- /dev/null +++ b/workflows/order_saga/tests/main_test.py @@ -0,0 +1,198 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +from typing import Any +from unittest import mock + +import flask +import pytest +import requests + +from .. import main + +app = flask.Flask(__name__) +ORDER_ID = "order-2001" +NOW = datetime.datetime.now(datetime.timezone.utc) + + +class FakeSnapshot: + def __init__(self, data: dict[str, Any] | None) -> None: + self._data = data + self.exists = data is not None + + def to_dict(self) -> dict[str, Any]: + return dict(self._data or {}) + + +class FakeDocument: + def __init__(self, store: dict[str, dict[str, Any]], doc_id: str) -> None: + self.store = store + self.id = doc_id + + def get(self, transaction: object = None) -> FakeSnapshot: + return FakeSnapshot(self.store.get(self.id)) + + +class FakeDatabase: + """An in-memory stand-in for the orders collection.""" + + def __init__(self) -> None: + self.store: dict[str, dict[str, Any]] = {} + + def collection(self, name: str) -> mock.MagicMock: + collection = mock.MagicMock() + collection.document.side_effect = lambda doc_id: FakeDocument( + self.store, doc_id + ) + return collection + + def transaction(self) -> mock.MagicMock: + # firestore.transactional begins, commits, and retries the transaction. + transaction = mock.MagicMock(_max_attempts=1, _id=b"test", _read_only=False) + transaction.update.side_effect = lambda ref, data: self.store[ref.id].update( + data + ) + return transaction + + +@pytest.fixture +def db(monkeypatch: pytest.MonkeyPatch) -> FakeDatabase: + database = FakeDatabase() + monkeypatch.setattr(main, "db", database) + return database + + +@pytest.fixture +def refund(monkeypatch: pytest.MonkeyPatch) -> mock.MagicMock: + refund_mock = mock.MagicMock() + monkeypatch.setattr(main, "refund", refund_mock) + return refund_mock + + +def create_order(db: FakeDatabase, **fields: object) -> dict[str, Any]: + db.store[ORDER_ID] = {"order_state": "paid", "amount_cents": 4999, **fields} + return db.store[ORDER_ID] + + +def call(body: dict[str, Any] | None = None) -> tuple[dict, int]: + body = {"order_id": ORDER_ID} if body is None else body + with app.test_request_context("/", method="POST", json=body): + return main.refund_order(flask.request) + + +def test_paid_order_is_refunded_once(db: FakeDatabase, refund: mock.MagicMock) -> None: + order = create_order(db) + + body, status = call() + + assert (status, body["already_processed"]) == (200, False) + refund.assert_called_once_with(ORDER_ID, 4999) + assert order["order_state"] == "refunded" + + +def test_refunded_order_is_not_refunded_again( + db: FakeDatabase, refund: mock.MagicMock +) -> None: + create_order(db, order_state="refunded") + + body, status = call() + + assert (status, body["already_processed"]) == (200, True) + refund.assert_not_called() + + +def test_order_with_active_claim_returns_503( + db: FakeDatabase, refund: mock.MagicMock +) -> None: + create_order( + db, + order_state="refund_in_progress", + attempt_id="other-attempt", + lease_expires=NOW + main.LEASE, + ) + + _, status = call() + + assert status == 503 + refund.assert_not_called() + + +@pytest.mark.parametrize( + "fields", + [ + { + "order_state": "refund_in_progress", + "attempt_id": "crashed-attempt", + "lease_expires": NOW - datetime.timedelta(seconds=1), + }, + {"order_state": "refund_failed"}, + ], + ids=["expired_claim", "failed_refund"], +) +def test_unfinished_refund_is_retried( + db: FakeDatabase, refund: mock.MagicMock, fields: dict[str, Any] +) -> None: + order = create_order(db, **fields) + + _, status = call() + + assert status == 200 + refund.assert_called_once() + assert order["order_state"] == "refunded" + + +def test_unpaid_order_is_not_refunded(db: FakeDatabase, refund: mock.MagicMock) -> None: + create_order(db, order_state="created") + + _, status = call() + + assert status == 409 + refund.assert_not_called() + + +def test_ambiguous_refund_failure_is_retryable( + db: FakeDatabase, refund: mock.MagicMock +) -> None: + order = create_order(db) + refund.side_effect = requests.Timeout("timed out") + + _, status = call() + + assert status == 503 + assert order["order_state"] == "refund_failed" + + +def test_stale_attempt_does_not_overwrite_newer_claim( + db: FakeDatabase, refund: mock.MagicMock +) -> None: + order = create_order(db) + refund.side_effect = lambda *args: order.update(attempt_id="newer-attempt") + + call() + + assert order["order_state"] == "refund_in_progress" + + +def test_refund_uses_its_own_idempotency_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(main, "PAYMENT_SERVICE_URL", "https://payments.example.com") + monkeypatch.setattr(main.id_token, "fetch_id_token", lambda *args: "token") + post = mock.MagicMock() + monkeypatch.setattr(main.requests, "post", post) + + main.refund(ORDER_ID, 4999) + + kwargs = post.call_args.kwargs + assert kwargs["headers"]["Idempotency-Key"] == f"refund-{ORDER_ID}" + assert kwargs["json"] == {"order_id": ORDER_ID, "amount_cents": -4999} diff --git a/workflows/order_saga/workflow.yaml b/workflows/order_saga/workflow.yaml new file mode 100644 index 00000000000..e1d35f92b92 --- /dev/null +++ b/workflows/order_saga/workflow.yaml @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START workflows_order_saga_workflow] +# Charges an order, then fulfills it. If fulfillment fails, refunds the charge, +# so the order never ends up paid but not fulfilled. Each step is safe to +# retry because each service processes an order ID only once. +# Deploy with the PAYMENT_FUNCTION_URL, FULFILLMENT_URL, and REFUND_FUNCTION_URL +# environment variables, and optionally FIRESTORE_DATABASE. +main: + params: [args] + steps: + - init: + assign: + - order_id: ${string(args.order_id)} + - database: ${default(sys.get_env("FIRESTORE_DATABASE"), "(default)")} + - documents: ${"projects/" + sys.get_env("GOOGLE_CLOUD_PROJECT_ID") + "/databases/" + database + "/documents"} + - create_order: + try: + call: googleapis.firestore.v1.projects.databases.documents.createDocument + args: + collectionId: orders + parent: ${documents} + documentId: ${order_id} + body: + fields: + order_state: + stringValue: created + amount_cents: + integerValue: ${string(args.amount_cents)} + except: + as: e + steps: + - check_conflict: + switch: + # A retried request continues with the existing order. + - condition: ${map.get(e, "code") != 409} + raise: ${e} + - charge: + try: + call: http.post + args: + url: ${sys.get_env("PAYMENT_FUNCTION_URL")} + auth: + type: OIDC + body: + order_id: ${order_id} + timeout: 30 + retry: ${http.default_retry} + - fulfill: + try: + call: http.post + args: + url: ${sys.get_env("FULFILLMENT_URL")} + auth: + type: OIDC + body: + order_id: ${order_id} + timeout: 30 + # Retry transient failures before treating fulfillment as failed. + retry: ${http.default_retry} + except: + as: e + steps: + - refund: + # The compensating step: undo the charge. + try: + call: http.post + args: + url: ${sys.get_env("REFUND_FUNCTION_URL")} + auth: + type: OIDC + body: + order_id: ${order_id} + timeout: 30 + retry: ${http.default_retry} + - report_refund: + return: + order_id: ${order_id} + status: refunded + reason: ${default(map.get(e, "message"), "fulfillment failed")} + - finish: + return: + order_id: ${order_id} + status: fulfilled +# [END workflows_order_saga_workflow]