From 8d697a8398d91aa860105083866a19a49ccf4c8b Mon Sep 17 00:00:00 2001
From: Matt Creaser <mattwcc@amazon.com>
Date: Tue, 14 May 2024 10:20:45 -0300
Subject: [PATCH] Add AuthMode to ModelSubscription

---
 .editorconfig                                 |   3 +-
 .idea/codeStyles/Project.xml                  |   1 -
 aws-api/api/aws-api.api                       |  15 ++
 .../api/aws/AppSyncGraphQLRequestFactory.kt   |  46 +++--
 .../api/graphql/model/ModelSubscription.kt    | 136 ++++++++++---
 .../graphql/model/ModelSubscriptionTest.kt    | 183 ++++++++++++++++--
 6 files changed, 325 insertions(+), 59 deletions(-)

diff --git a/.editorconfig b/.editorconfig
index 7b8a8aa853..a37bf220da 100644
--- a/.editorconfig
+++ b/.editorconfig
@@ -1,3 +1,4 @@
 [*.{kt,kts}]
 #this is to match java checkstyle
-max_line_length=120
\ No newline at end of file
+max_line_length=120
+ktlint_code_style=android_studio
diff --git a/.idea/codeStyles/Project.xml b/.idea/codeStyles/Project.xml
index 80a2c0abf4..b41da14643 100644
--- a/.idea/codeStyles/Project.xml
+++ b/.idea/codeStyles/Project.xml
@@ -1,6 +1,5 @@
 <component name="ProjectCodeStyleConfiguration">
   <code_scheme name="Project" version="173">
-    <option name="ENABLE_SECOND_REFORMAT" value="true" />
     <JavaCodeStyleSettings>
       <option name="IMPORT_LAYOUT_TABLE">
         <value>
diff --git a/aws-api/api/aws-api.api b/aws-api/api/aws-api.api
index bfd34be952..521a9d551f 100644
--- a/aws-api/api/aws-api.api
+++ b/aws-api/api/aws-api.api
@@ -83,7 +83,10 @@ public final class com/amplifyframework/api/aws/AppSyncGraphQLRequestFactory {
 	public static final fun buildQuery (Ljava/lang/Class;Ljava/lang/String;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun buildQuery (Ljava/lang/Class;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun buildSubscription (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun buildSubscription (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun buildSubscription (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun buildSubscription (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static synthetic fun buildSubscription$default (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;ILjava/lang/Object;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 }
 
 public final class com/amplifyframework/api/aws/BuildConfig {
@@ -263,12 +266,24 @@ public final class com/amplifyframework/api/graphql/model/ModelQuery {
 public final class com/amplifyframework/api/graphql/model/ModelSubscription {
 	public static final field INSTANCE Lcom/amplifyframework/api/graphql/model/ModelSubscription;
 	public static final fun of (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun of (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun of (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun of (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static synthetic fun of$default (Ljava/lang/Class;Lcom/amplifyframework/api/graphql/SubscriptionType;Lcom/amplifyframework/api/aws/AuthorizationType;ILjava/lang/Object;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onCreate (Ljava/lang/Class;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onCreate (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onCreate (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onCreate (Ljava/lang/Class;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static synthetic fun onCreate$default (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;ILjava/lang/Object;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onDelete (Ljava/lang/Class;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onDelete (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onDelete (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onDelete (Ljava/lang/Class;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static synthetic fun onDelete$default (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;ILjava/lang/Object;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onUpdate (Ljava/lang/Class;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onUpdate (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static final fun onUpdate (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 	public static final fun onUpdate (Ljava/lang/Class;Lkotlin/jvm/functions/Function1;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
+	public static synthetic fun onUpdate$default (Ljava/lang/Class;Lcom/amplifyframework/api/aws/AuthorizationType;ILjava/lang/Object;)Lcom/amplifyframework/api/graphql/GraphQLRequest;
 }
 
diff --git a/aws-api/src/main/java/com/amplifyframework/api/aws/AppSyncGraphQLRequestFactory.kt b/aws-api/src/main/java/com/amplifyframework/api/aws/AppSyncGraphQLRequestFactory.kt
index b5faa3b7ed..aec1db6fd6 100644
--- a/aws-api/src/main/java/com/amplifyframework/api/aws/AppSyncGraphQLRequestFactory.kt
+++ b/aws-api/src/main/java/com/amplifyframework/api/aws/AppSyncGraphQLRequestFactory.kt
@@ -482,18 +482,19 @@ object AppSyncGraphQLRequestFactory {
      * Creates a [GraphQLRequest] that represents a subscription of a given type.
      * @param modelClass the model type.
      * @param subscriptionType the subscription type.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param <R> the response type.
      * @param <T> the concrete model type.
      * @return a valid [GraphQLRequest] instance.
      * @throws IllegalStateException when the model schema does not contain the expected information.
-     </T></R> */
+     */
     @JvmStatic
+    @JvmOverloads
     fun <R, T : Model> buildSubscription(
         modelClass: Class<T>,
-        subscriptionType: SubscriptionType
-    ): GraphQLRequest<R> {
-        return buildSubscriptionInternal(modelClass, subscriptionType, null)
-    }
+        subscriptionType: SubscriptionType,
+        authMode: AuthorizationType? = null
+    ) = buildSubscriptionInternal<R, T, ModelPath<T>>(modelClass, subscriptionType, authMode, null)
 
     /**
      * Creates a [GraphQLRequest] that represents a subscription of a given type.
@@ -505,19 +506,38 @@ object AppSyncGraphQLRequestFactory {
      * @param <P> the concrete model path for the M model type
      * @return a valid [GraphQLRequest] instance.
      * @throws IllegalStateException when the model schema does not contain the expected information.
-     </T></R> */
+     */
     @JvmStatic
     fun <R, T : Model, P : ModelPath<T>> buildSubscription(
         modelClass: Class<T>,
         subscriptionType: SubscriptionType,
         includes: ((P) -> List<PropertyContainerPath>)
-    ): GraphQLRequest<R> {
-        return buildSubscriptionInternal(modelClass, subscriptionType, includes)
-    }
+    ) = buildSubscriptionInternal<R, T, P>(modelClass, subscriptionType, null, includes)
+
+    /**
+     * Creates a [GraphQLRequest] that represents a subscription of a given type.
+     * @param modelClass the model type.
+     * @param subscriptionType the subscription type.
+     * @param authMode The [AuthorizationType] to use for making the request
+     * @param includes lambda returning list of relationships that should be included in the selection set
+     * @param <R> the response type.
+     * @param <T> the concrete model type.
+     * @param <P> the concrete model path for the M model type
+     * @return a valid [GraphQLRequest] instance.
+     * @throws IllegalStateException when the model schema does not contain the expected information.
+     */
+    @JvmStatic
+    fun <R, T : Model, P : ModelPath<T>> buildSubscription(
+        modelClass: Class<T>,
+        subscriptionType: SubscriptionType,
+        authMode: AuthorizationType,
+        includes: ((P) -> List<PropertyContainerPath>)
+    ) = buildSubscriptionInternal<R, T, P>(modelClass, subscriptionType, authMode, includes)
 
     private fun <R, T : Model, P : ModelPath<T>> buildSubscriptionInternal(
         modelClass: Class<T>,
         subscriptionType: SubscriptionType,
+        authMode: AuthorizationType?,
         includes: ((P) -> List<PropertyContainerPath>)?
     ): GraphQLRequest<R> {
         return try {
@@ -527,8 +547,12 @@ object AppSyncGraphQLRequestFactory {
                 .requestOptions(ApiGraphQLRequestOptions())
                 .responseType(modelClass)
 
-            val customSelectionSet = includes?.let { createApiSelectionSet(modelClass, subscriptionType, it) }
-            customSelectionSet?.let { builder.selectionSet(it) }
+            authMode?.let { builder.authorizationType(it) }
+
+            includes?.let {
+                val customSelectionSet = createApiSelectionSet(modelClass, subscriptionType, it)
+                builder.selectionSet(customSelectionSet)
+            }
 
             builder.build()
         } catch (exception: AmplifyException) {
diff --git a/aws-api/src/main/java/com/amplifyframework/api/graphql/model/ModelSubscription.kt b/aws-api/src/main/java/com/amplifyframework/api/graphql/model/ModelSubscription.kt
index b897bf4b7f..920bc2cfc1 100644
--- a/aws-api/src/main/java/com/amplifyframework/api/graphql/model/ModelSubscription.kt
+++ b/aws-api/src/main/java/com/amplifyframework/api/graphql/model/ModelSubscription.kt
@@ -15,6 +15,7 @@
 package com.amplifyframework.api.graphql.model
 
 import com.amplifyframework.api.aws.AppSyncGraphQLRequestFactory.buildSubscription
+import com.amplifyframework.api.aws.AuthorizationType
 import com.amplifyframework.api.graphql.GraphQLRequest
 import com.amplifyframework.api.graphql.SubscriptionType
 import com.amplifyframework.core.model.Model
@@ -30,16 +31,17 @@ object ModelSubscription {
      * Builds a subscriptions request of a given `type` for a `modelType`.
      * @param modelType the model class.
      * @param type the subscription type.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param <M> the concrete type of the model.
      * @return a valid [GraphQLRequest] instance.
-     </M> */
+     */
     @JvmStatic
+    @JvmOverloads
     fun <M : Model> of(
         modelType: Class<M>,
         type: SubscriptionType,
-    ): GraphQLRequest<M> {
-        return buildSubscription(modelType, type)
-    }
+        authMode: AuthorizationType? = null
+    ): GraphQLRequest<M> = buildSubscription(modelType, type, authMode)
 
     /**
      * Builds a subscriptions request of a given `type` for a `modelType`.
@@ -49,100 +51,170 @@ object ModelSubscription {
      * @param <M> the concrete type of the model.
      * @param <P> the concrete model path for the M model type
      * @return a valid [GraphQLRequest] instance.
-     </M> */
+     */
     @JvmStatic
     fun <M : Model, P : ModelPath<M>> of(
         modelType: Class<M>,
         type: SubscriptionType,
         includes: ((P) -> List<PropertyContainerPath>)
-    ): GraphQLRequest<M> {
-        return buildSubscription(modelType, type, includes)
-    }
+    ): GraphQLRequest<M> = buildSubscription(modelType, type, includes)
+
+    /**
+     * Builds a subscriptions request of a given `type` for a `modelType`.
+     * @param modelType the model class.
+     * @param type the subscription type.
+     * @param authMode The [AuthorizationType] to use for making the request
+     * @param includes lambda returning list of relationships that should be included in the selection set
+     * @param <M> the concrete type of the model.
+     * @param <P> the concrete model path for the M model type
+     * @return a valid [GraphQLRequest] instance.
+     */
+    @JvmStatic
+    fun <M : Model, P : ModelPath<M>> of(
+        modelType: Class<M>,
+        type: SubscriptionType,
+        authMode: AuthorizationType,
+        includes: ((P) -> List<PropertyContainerPath>)
+    ): GraphQLRequest<M> = buildSubscription(modelType, type, authMode, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_CREATE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param <M> the concrete type of the model.
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
-    fun <M : Model> onCreate(modelType: Class<M>): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_CREATE)
-    }
+    @JvmOverloads
+    fun <M : Model> onCreate(
+        modelType: Class<M>,
+        authMode: AuthorizationType? = null
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_CREATE, authMode)
+
+    /**
+     * Creates a subscription request of type [SubscriptionType.ON_CREATE].
+     * @param modelType the model class.
+     * @param includes lambda returning list of relationships that should be included in the selection set
+     * @param <M> the concrete type of the model.
+     * @param <P> the concrete model path for the M model type
+     * @return a valid [GraphQLRequest] instance.
+     * @see .of
+     */
+    @JvmStatic
+    fun <M : Model, P : ModelPath<M>> onCreate(
+        modelType: Class<M>,
+        includes: ((P) -> List<PropertyContainerPath>)
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_CREATE, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_CREATE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param includes lambda returning list of relationships that should be included in the selection set
      * @param <M> the concrete type of the model.
      * @param <P> the concrete model path for the M model type
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
     fun <M : Model, P : ModelPath<M>> onCreate(
         modelType: Class<M>,
+        authMode: AuthorizationType,
         includes: ((P) -> List<PropertyContainerPath>)
-    ): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_CREATE, includes)
-    }
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_CREATE, authMode, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_DELETE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param <M> the concrete type of the model.
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
-    fun <M : Model> onDelete(modelType: Class<M>): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_DELETE)
-    }
+    @JvmOverloads
+    fun <M : Model> onDelete(
+        modelType: Class<M>,
+        authMode: AuthorizationType? = null
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_DELETE, authMode)
+
+    /**
+     * Creates a subscription request of type [SubscriptionType.ON_DELETE].
+     * @param modelType the model class.
+     * @param includes lambda returning list of relationships that should be included in the selection set
+     * @param <M> the concrete type of the model.
+     * @param <P> the concrete model path for the M model type
+     * @return a valid [GraphQLRequest] instance.
+     * @see .of
+     */
+    @JvmStatic
+    fun <M : Model, P : ModelPath<M>> onDelete(
+        modelType: Class<M>,
+        includes: ((P) -> List<PropertyContainerPath>)
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_DELETE, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_DELETE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param includes lambda returning list of relationships that should be included in the selection set
      * @param <M> the concrete type of the model.
      * @param <P> the concrete model path for the M model type
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
     fun <M : Model, P : ModelPath<M>> onDelete(
         modelType: Class<M>,
+        authMode: AuthorizationType,
         includes: ((P) -> List<PropertyContainerPath>)
-    ): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_DELETE, includes)
-    }
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_DELETE, authMode, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_UPDATE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param <M> the concrete type of the model.
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
-    fun <M : Model> onUpdate(modelType: Class<M>): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_UPDATE)
-    }
+    @JvmOverloads
+    fun <M : Model> onUpdate(
+        modelType: Class<M>,
+        authMode: AuthorizationType? = null
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_UPDATE, authMode)
+
+    /**
+     * Creates a subscription request of type [SubscriptionType.ON_UPDATE].
+     * @param modelType the model class.
+     * @param includes lambda returning list of relationships that should be included in the selection set
+     * @param <M> the concrete type of the model.
+     * @param <P> the concrete model path for the M model type
+     * @return a valid [GraphQLRequest] instance.
+     * @see .of
+     */
+    @JvmStatic
+    fun <M : Model, P : ModelPath<M>> onUpdate(
+        modelType: Class<M>,
+        includes: ((P) -> List<PropertyContainerPath>)
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_UPDATE, includes)
 
     /**
      * Creates a subscription request of type [SubscriptionType.ON_UPDATE].
      * @param modelType the model class.
+     * @param authMode The [AuthorizationType] to use for making the request
      * @param includes lambda returning list of relationships that should be included in the selection set
      * @param <M> the concrete type of the model.
      * @param <P> the concrete model path for the M model type
      * @return a valid [GraphQLRequest] instance.
      * @see .of
-     </M> */
+     */
     @JvmStatic
     fun <M : Model, P : ModelPath<M>> onUpdate(
         modelType: Class<M>,
+        authMode: AuthorizationType,
         includes: ((P) -> List<PropertyContainerPath>)
-    ): GraphQLRequest<M> {
-        return of(modelType, SubscriptionType.ON_UPDATE, includes)
-    }
+    ): GraphQLRequest<M> = of(modelType, SubscriptionType.ON_UPDATE, authMode, includes)
 }
diff --git a/aws-api/src/test/java/com/amplifyframework/graphql/model/ModelSubscriptionTest.kt b/aws-api/src/test/java/com/amplifyframework/graphql/model/ModelSubscriptionTest.kt
index 4e3a00f0ee..1e9659032a 100644
--- a/aws-api/src/test/java/com/amplifyframework/graphql/model/ModelSubscriptionTest.kt
+++ b/aws-api/src/test/java/com/amplifyframework/graphql/model/ModelSubscriptionTest.kt
@@ -16,12 +16,14 @@
 package com.amplifyframework.graphql.model
 
 import com.amplifyframework.api.aws.AppSyncGraphQLRequestFactory
+import com.amplifyframework.api.aws.AuthorizationType
 import com.amplifyframework.api.graphql.GraphQLRequest
 import com.amplifyframework.api.graphql.SubscriptionType
 import com.amplifyframework.api.graphql.model.ModelSubscription
 import com.amplifyframework.core.model.includes
 import com.amplifyframework.testmodels.lazy.Post
 import com.amplifyframework.testmodels.lazy.PostPath
+import io.kotest.matchers.shouldBe
 import org.junit.Assert.assertEquals
 import org.junit.Test
 
@@ -42,6 +44,23 @@ class ModelSubscriptionTest {
         assertEquals(expectedRequest, actualRequest)
     }
 
+    @Test
+    fun `of with authMode`() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_CREATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType,
+            expectedAuthMode
+        )
+
+        val actualRequest = ModelSubscription.of(expectedClass, expectedType, expectedAuthMode)
+
+        actualRequest shouldBe expectedRequest
+    }
+
     @Test
     fun of_with_includes() {
         val expectedClass = Post::class.java
@@ -63,21 +82,59 @@ class ModelSubscriptionTest {
     }
 
     @Test
-    fun create() {
+    fun `of with includes and authMode`() {
         val expectedClass = Post::class.java
         val expectedType = SubscriptionType.ON_CREATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
 
         val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory
-            .buildSubscription(
+            .buildSubscription<Post, Post, PostPath>(
                 expectedClass,
-                expectedType
-            )
+                expectedType,
+                expectedAuthMode
+            ) {
+                includes(it.comments, it.blog)
+            }
+
+        val actualRequest = ModelSubscription.of<Post, PostPath>(expectedClass, expectedType, expectedAuthMode) {
+            includes(it.comments, it.blog)
+        }
+
+        actualRequest shouldBe expectedRequest
+    }
+
+    @Test
+    fun create() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_CREATE
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType
+        )
 
         val actualRequest = ModelSubscription.onCreate(expectedClass)
 
         assertEquals(expectedRequest, actualRequest)
     }
 
+    @Test
+    fun `create with authMode`() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_CREATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType,
+            expectedAuthMode
+        )
+
+        val actualRequest = ModelSubscription.onCreate(expectedClass, expectedAuthMode)
+
+        actualRequest shouldBe expectedRequest
+    }
+
     @Test
     fun create_with_includes() {
         val expectedClass = Post::class.java
@@ -99,21 +156,59 @@ class ModelSubscriptionTest {
     }
 
     @Test
-    fun delete() {
+    fun `create with includes and authMode`() {
         val expectedClass = Post::class.java
-        val expectedType = SubscriptionType.ON_DELETE
+        val expectedType = SubscriptionType.ON_CREATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
 
         val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory
-            .buildSubscription(
+            .buildSubscription<Post, Post, PostPath>(
                 expectedClass,
-                expectedType
-            )
+                expectedType,
+                expectedAuthMode
+            ) {
+                includes(it.comments, it.blog)
+            }
+
+        val actualRequest = ModelSubscription.onCreate<Post, PostPath>(expectedClass, expectedAuthMode) {
+            includes(it.comments, it.blog)
+        }
+
+        actualRequest shouldBe expectedRequest
+    }
+
+    @Test
+    fun delete() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_DELETE
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType
+        )
 
         val actualRequest = ModelSubscription.onDelete(expectedClass)
 
         assertEquals(expectedRequest, actualRequest)
     }
 
+    @Test
+    fun `delete with authMode`() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_DELETE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType,
+            expectedAuthMode
+        )
+
+        val actualRequest = ModelSubscription.onDelete(expectedClass, expectedAuthMode)
+
+        actualRequest shouldBe expectedRequest
+    }
+
     @Test
     fun delete_with_includes() {
         val expectedClass = Post::class.java
@@ -135,21 +230,59 @@ class ModelSubscriptionTest {
     }
 
     @Test
-    fun update() {
+    fun `delete with includes and authMode`() {
         val expectedClass = Post::class.java
-        val expectedType = SubscriptionType.ON_UPDATE
+        val expectedType = SubscriptionType.ON_DELETE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
 
         val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory
-            .buildSubscription(
+            .buildSubscription<Post, Post, PostPath>(
                 expectedClass,
-                expectedType
-            )
+                expectedType,
+                expectedAuthMode
+            ) {
+                includes(it.comments, it.blog)
+            }
+
+        val actualRequest = ModelSubscription.onDelete<Post, PostPath>(expectedClass, expectedAuthMode) {
+            includes(it.comments, it.blog)
+        }
+
+        actualRequest shouldBe expectedRequest
+    }
+
+    @Test
+    fun update() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_UPDATE
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType
+        )
 
         val actualRequest = ModelSubscription.onUpdate(expectedClass)
 
         assertEquals(expectedRequest, actualRequest)
     }
 
+    @Test
+    fun `update with authMode`() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_UPDATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory.buildSubscription(
+            expectedClass,
+            expectedType,
+            expectedAuthMode
+        )
+
+        val actualRequest = ModelSubscription.onUpdate(expectedClass, authMode = expectedAuthMode)
+
+        actualRequest shouldBe expectedRequest
+    }
+
     @Test
     fun update_with_includes() {
         val expectedClass = Post::class.java
@@ -169,4 +302,26 @@ class ModelSubscriptionTest {
 
         assertEquals(expectedRequest, actualRequest)
     }
+
+    @Test
+    fun `update with includes and authMode`() {
+        val expectedClass = Post::class.java
+        val expectedType = SubscriptionType.ON_UPDATE
+        val expectedAuthMode = AuthorizationType.OPENID_CONNECT
+
+        val expectedRequest: GraphQLRequest<Post> = AppSyncGraphQLRequestFactory
+            .buildSubscription<Post, Post, PostPath>(
+                expectedClass,
+                expectedType,
+                expectedAuthMode
+            ) {
+                includes(it.comments, it.blog)
+            }
+
+        val actualRequest = ModelSubscription.onUpdate<Post, PostPath>(expectedClass, expectedAuthMode) {
+            includes(it.comments, it.blog)
+        }
+
+        actualRequest shouldBe expectedRequest
+    }
 }