diff --git a/firebase-appdistribution/test-app/src/main/java/com/googletest/firebase/appdistribution/testapp/MainActivity.kt b/firebase-appdistribution/test-app/src/main/java/com/googletest/firebase/appdistribution/testapp/MainActivity.kt
index 8ebb6028189..19c9721229c 100644
--- a/firebase-appdistribution/test-app/src/main/java/com/googletest/firebase/appdistribution/testapp/MainActivity.kt
+++ b/firebase-appdistribution/test-app/src/main/java/com/googletest/firebase/appdistribution/testapp/MainActivity.kt
@@ -12,13 +12,14 @@ import androidx.appcompat.app.AppCompatActivity
 import androidx.appcompat.widget.AppCompatButton
 import com.google.android.gms.tasks.Task
 import com.google.firebase.appdistribution.AppDistributionRelease
-import com.google.firebase.appdistribution.FirebaseAppDistribution
 import com.google.firebase.appdistribution.UpdateProgress
+import com.google.firebase.appdistribution.ktx.appDistribution
+import com.google.firebase.ktx.Firebase
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
 
 class MainActivity : AppCompatActivity() {
-    var firebaseAppDistribution: FirebaseAppDistribution = FirebaseAppDistribution.getInstance()
+    var firebaseAppDistribution = Firebase.appDistribution
     var updateTask: Task<Void>? = null
     var release: AppDistributionRelease? = null
     val executorService: ExecutorService = Executors.newFixedThreadPool(1)
diff --git a/firebase-appdistribution/test-app/test-app.gradle b/firebase-appdistribution/test-app/test-app.gradle
index 4350447264d..137a53d6b14 100644
--- a/firebase-appdistribution/test-app/test-app.gradle
+++ b/firebase-appdistribution/test-app/test-app.gradle
@@ -19,12 +19,12 @@ plugins {
 }
 
 android {
-    compileSdkVersion 30
+    compileSdkVersion 31
 
     defaultConfig {
         applicationId "com.googletest.firebase.appdistribution.testapp"
         minSdkVersion 19
-        targetSdkVersion 30
+        targetSdkVersion 31
         versionName "1.0"
         versionCode 1
 
@@ -33,7 +33,7 @@ android {
 
     buildTypes {
         release {
-            // This is so we can build the "release" variant for the "dev-app" (and the SDK) without
+            // This is so we can build the "release" variant for the "test-app" (and the SDK) without
             // needing to have the app signed.
             initWith debug
         }
@@ -45,6 +45,7 @@ android {
             }
         }
     }
+
     compileOptions {
         sourceCompatibility JavaVersion.VERSION_1_8
         targetCompatibility JavaVersion.VERSION_1_8
@@ -58,13 +59,18 @@ dependencies {
     // TODO(rachelprince): Add flag to build with public version of SDK
     println("Building with HEAD version ':firebase-appdistribution' of SDK")
 
-    // Include the API in all variants
-    implementation project(':firebase-appdistribution-api')
-    // Only include the full SDK implementation in the debug variant
-    debugImplementation project(':firebase-appdistribution')
+    // All variants use the API
+    implementation project(':firebase-appdistribution-api:ktx')
 
+    // In this test project we also need to explicitly declare these dependencies
+    implementation project(':firebase-appdistribution-api')
+    implementation project(':firebase-common:ktx')
     implementation "com.google.android.gms:play-services-tasks:18.0.1"
 
+    // Debug uses the full implementation
+    debugImplementation project(':firebase-appdistribution')
+
+    // Other dependencies
     implementation "org.jetbrains.kotlin:kotlin-stdlib:$kotlinVersion"
     implementation 'androidx.core:core-ktx:1.5.0'
     implementation "androidx.core:core:1.6.0"