From af0ea5d98fa1858fb973e2021e9868feabe501ed Mon Sep 17 00:00:00 2001
From: Matt Creaser <mattwcc@amazon.com>
Date: Wed, 15 May 2024 10:44:43 -0300
Subject: [PATCH] Fix the handling of exceptions when subscribing with Kotlin
 Facade

---
 core-kotlin/build.gradle.kts                  |  1 +
 .../kotlin/api/KotlinApiFacade.kt             |  3 ++-
 .../kotlin/api/KotlinApiFacadeTest.kt         | 21 +++++++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/core-kotlin/build.gradle.kts b/core-kotlin/build.gradle.kts
index d835cecf41..681f069a9b 100644
--- a/core-kotlin/build.gradle.kts
+++ b/core-kotlin/build.gradle.kts
@@ -33,6 +33,7 @@ dependencies {
     testImplementation(libs.test.mockk)
     testImplementation(libs.test.kotlin.coroutines)
     testImplementation(project(":testmodels"))
+    testImplementation(libs.test.kotest.assertions)
 }
 
 android.kotlinOptions {
diff --git a/core-kotlin/src/main/java/com/amplifyframework/kotlin/api/KotlinApiFacade.kt b/core-kotlin/src/main/java/com/amplifyframework/kotlin/api/KotlinApiFacade.kt
index 722b9fd769..b4ece4439d 100644
--- a/core-kotlin/src/main/java/com/amplifyframework/kotlin/api/KotlinApiFacade.kt
+++ b/core-kotlin/src/main/java/com/amplifyframework/kotlin/api/KotlinApiFacade.kt
@@ -108,7 +108,8 @@ class KotlinApiFacade(private val delegate: Delegate = Amplify.API) : Api {
                 { subscription.completions.tryEmit(Unit) }
             )
         }
-        subscription.cancelable = operation as Cancelable
+        // If subscribe fails it does not return an operation, and instead invokes the onSubscriptionFailure callback
+        operation?.let { subscription.cancelable = operation }
         return subscription.awaitStart()
     }
 
diff --git a/core-kotlin/src/test/java/com/amplifyframework/kotlin/api/KotlinApiFacadeTest.kt b/core-kotlin/src/test/java/com/amplifyframework/kotlin/api/KotlinApiFacadeTest.kt
index ca2bdf80ee..3c07c4358a 100644
--- a/core-kotlin/src/test/java/com/amplifyframework/kotlin/api/KotlinApiFacadeTest.kt
+++ b/core-kotlin/src/test/java/com/amplifyframework/kotlin/api/KotlinApiFacadeTest.kt
@@ -23,6 +23,7 @@ import com.amplifyframework.api.graphql.GraphQLResponse
 import com.amplifyframework.api.rest.RestOptions
 import com.amplifyframework.api.rest.RestResponse
 import com.amplifyframework.core.Consumer
+import io.kotest.assertions.throwables.shouldThrow
 import io.mockk.every
 import io.mockk.mockk
 import kotlinx.coroutines.Dispatchers
@@ -32,6 +33,7 @@ import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
 import org.junit.Test
 
@@ -174,6 +176,25 @@ class KotlinApiFacadeTest {
         api.subscribe(request).first()
     }
 
+    @Test
+    fun `subscribe throws when exception occurs during subscription establishment`() = runTest {
+        val request = mockk<GraphQLRequest<String>>()
+        val expectedFailure = ApiException("uh", "oh")
+
+        every {
+            delegate.subscribe(eq(request), any(), any(), any(), any())
+        } answers {
+            val indexOfErrorConsumer = 3
+            val onError = it.invocation.args[indexOfErrorConsumer] as Consumer<ApiException>
+            GlobalScope.launch(Dispatchers.IO) { onError.accept(expectedFailure) }
+            null
+        }
+
+        shouldThrow<ApiException> {
+            api.subscribe(request).first()
+        }
+    }
+
     /**
      * When the underlying get() emits a response,
      * it should be returned from the coroutine API.