From e6dbb34fd6d68004af5d2d790c4cb5a64100943e Mon Sep 17 00:00:00 2001
From: jorgee <jorge.ejarque@seqera.io>
Date: Fri, 28 Feb 2025 16:49:21 +0100
Subject: [PATCH] Fix Google Batch autoRetryExitCodes bug

Signed-off-by: jorgee <jorge.ejarque@seqera.io>
---
 .../nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy | 4 +++-
 .../cloud/google/batch/GoogleBatchTaskHandlerTest.groovy      | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
index 442fea9a6a..f9f467a784 100644
--- a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
+++ b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
@@ -17,6 +17,8 @@
 
 package nextflow.cloud.google.batch
 
+import nextflow.cloud.google.batch.client.BatchConfig
+
 import java.nio.file.Path
 import java.util.regex.Pattern
 
@@ -268,7 +270,7 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
                     LifecyclePolicy.newBuilder()
                         .setActionCondition(
                             LifecyclePolicy.ActionCondition.newBuilder()
-                                .addExitCodes(50001)
+                                .addAllExitCodes(executor.config.autoRetryExitCodes)
                         )
                         .setAction(LifecyclePolicy.Action.RETRY_TASK)
                 )
diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
index 37b27e0b5a..be9a6b0bb0 100644
--- a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
+++ b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy
@@ -146,6 +146,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
                 getBootDiskImage() >> BOOT_IMAGE
                 getCpuPlatform() >> CPU_PLATFORM
                 getMaxSpotAttempts() >> 5
+                getAutoRetryExitCodes() >> [50001,50002]
                 getSpot() >> true
                 getNetwork() >> 'net-1'
                 getServiceAccountEmail() >> 'foo@bar.baz'
@@ -198,7 +199,9 @@ class GoogleBatchTaskHandlerTest extends Specification {
         taskSpec.getMaxRunDuration().getSeconds() == TIMEOUT.seconds
         taskSpec.getVolumes(0).getMountPath() == '/tmp'
         taskSpec.getMaxRetryCount() == 5
+        taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodesCount() == 2
         taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodes(0) == 50001
+        taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodes(1) == 50002
         taskSpec.getLifecyclePolicies(0).getAction().toString() == 'RETRY_TASK'
         and:
         runnable.getContainer().getCommandsList().join(' ') == '/bin/bash -o pipefail -c bash .command.run'