From e67211d38139a8f9d63db055b82a6bff93b2c5bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Ga=C5=82uszka?= <galuszkak@gmail.com>
Date: Sun, 24 Aug 2025 14:34:44 +0000
Subject: [PATCH] feat: add support for --wait flag in jobs

---
 action.yml              |  7 +++++++
 src/main.ts             |  5 +++++
 tests/unit/main.test.ts | 24 ++++++++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/action.yml b/action.yml
index 6ef53868..3eadd49a 100644
--- a/action.yml
+++ b/action.yml
@@ -229,6 +229,13 @@ inputs:
     default: 'false'
     required: false
 
+  wait:
+    description: |-
+      If true, the action will wait for the job to complete before exiting. This
+      option only applies to jobs.
+    default: 'true'
+    required: false
+
   revision_traffic:
     description: |-
       Comma-separated list of revision traffic assignments.
diff --git a/src/main.ts b/src/main.ts
index a18900c9..cd3b1caa 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -113,6 +113,7 @@ export async function run(): Promise<void> {
     const tag = getInput('tag');
     const timeout = getInput('timeout');
     const noTraffic = (getInput('no_traffic') || '').toLowerCase() === 'true';
+    const wait = parseBoolean(getInput('wait'));
     const revTraffic = getInput('revision_traffic');
     const tagTraffic = getInput('tag_traffic');
     const labels = parseKVString(getInput('labels'));
@@ -196,6 +197,10 @@ export async function run(): Promise<void> {
       setEnvVarsFlags(deployCmd, envVars, envVarsFile, envVarsUpdateStrategy);
       setSecretsFlags(deployCmd, secrets, secretsUpdateStrategy);
 
+      if (wait) {
+        deployCmd.push('--wait');
+      }
+
       // There is no --update-secrets flag on jobs, but there will be in the
       // future. At that point, we can remove this.
       const idx = deployCmd.indexOf('--update-secrets');
diff --git a/tests/unit/main.test.ts b/tests/unit/main.test.ts
index 82f13dc2..882d43ee 100644
--- a/tests/unit/main.test.ts
+++ b/tests/unit/main.test.ts
@@ -520,6 +520,30 @@ test('#run', { concurrency: true }, async (suite) => {
     const args = mocks.getExecOutput.mock.calls?.at(0)?.arguments?.at(1);
     assertMembers(args, ['run', 'jobs', 'deploy', 'my-test-job']);
   });
+
+  await suite.test('deploys a job with --wait', async (t) => {
+    const mocks = defaultMocks(t.mock, {
+      job: 'my-test-job',
+      wait: 'true',
+    });
+
+    await run();
+
+    const args = mocks.getExecOutput.mock.calls?.at(0)?.arguments?.at(1);
+    assert.ok(args?.includes('--wait'));
+  });
+
+  await suite.test('deploys a job without --wait', async (t) => {
+    const mocks = defaultMocks(t.mock, {
+      job: 'my-test-job',
+      wait: 'false',
+    });
+
+    await run();
+
+    const args = mocks.getExecOutput.mock.calls?.at(0)?.arguments?.at(1);
+    assert.ok(!args.includes('--wait'));
+  });
 });
 
 const splitKV = (s: string): Record<string, string> => {