Handle GoogleAI 503 errors in test retry logic (#211)

Author: aozherelyeva

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/MultipleSystemMessagesPromptIntegrationTest.kt

@@ -1,5 +1,6 @@
 package ai.koog.integration.tests
 
+import ai.koog.integration.tests.utils.TestUtils.executeWithRetry
 import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
@@ -45,9 +46,9 @@ class MultipleSystemMessagesPromptIntegrationTest {
         val modelAnthropic = AnthropicModels.Haiku_3_5
         val modelGemini = GoogleModels.Gemini2_0Flash
 
-        val responseOpenAI = executor.execute(prompt, modelOpenAI)
-        val responseAnthropic = executor.execute(prompt, modelAnthropic)
-        val responseGemini = executor.execute(prompt, modelGemini)
+        val responseOpenAI = executeWithRetry { executor.execute(prompt, modelOpenAI) }
+        val responseAnthropic = executeWithRetry { executor.execute(prompt, modelAnthropic) }
+        val responseGemini = executeWithRetry { executor.execute(prompt, modelGemini) }
 
         assertTrue(responseOpenAI.content.isNotEmpty(), "OpenAI response should not be empty")
         assertTrue(responseAnthropic.content.isNotEmpty(), "Anthropic response should not be empty")

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/TestUtils.kt

@@ -29,6 +29,7 @@ internal object TestUtils {
 
     private const val GOOGLE_API_ERROR = "Field 'parts' is required for type with serial name"
     private const val GOOGLE_500_ERROR = "Error from GoogleAI API: 500 Internal Server Error"
+    private const val GOOGLE_503_ERROR = "Error from GoogleAI API: 503 Service Unavailable"
 
     suspend fun <T> executeWithRetry(operation: suspend () -> T): T {
         val maxRetries = 3
@@ -39,8 +40,8 @@ internal object TestUtils {
             try {
                 return operation()
             } catch (e: Exception) {
-                if (e.message?.contains(GOOGLE_500_ERROR) == true) {
-                    assumeTrue(false, "Skipping test due to GoogleAI API 500 Internal Server Error")
+                if (e.message?.contains(GOOGLE_500_ERROR) == true || e.message?.contains(GOOGLE_503_ERROR) == true) {
+                    assumeTrue(false, "Skipping test due to ${e.message}")
                 } else if (e.message?.contains(GOOGLE_API_ERROR) == true) {
                     lastException = e
                     attempts++
@@ -69,8 +70,8 @@ internal object TestUtils {
             try {
                 return operation()
             } catch (e: Exception) {
-                if (e.message?.contains(GOOGLE_500_ERROR) == true) {
-                    assumeTrue(false, "Skipping test due to GoogleAI API 500 Internal Server Error")
+                if (e.message?.contains(GOOGLE_500_ERROR) == true || e.message?.contains(GOOGLE_503_ERROR) == true) {
+                    assumeTrue(false, "Skipping test due to ${e.message}}")
                 } else {
                     attempts++
                     println("Attempt $attempts/$maxRetries failed with exception ${e.message}, retrying...")