diff --git a/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModel.kt b/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModel.kt
index 3fd4d51e472..95aee4bacfc 100644
--- a/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModel.kt
+++ b/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModel.kt
@@ -64,11 +64,8 @@ class AutoCredentialViewModel(
      */
     interface ChallengeGenerator {
 
-        /** Get callback that will be called later after challenge generated */
-        fun getCallback(): GenerateChallengeCallback?
-
-        /** Set callback that will be called later after challenge generated */
-        fun setCallback(callback: GenerateChallengeCallback?)
+        /** Callback that will be called later after challenge generated */
+        var callback: GenerateChallengeCallback?
 
         /** Method for generating challenge from FingerprintManager or FaceManager */
         fun generateChallenge(userId: Int)
@@ -79,30 +76,16 @@ class AutoCredentialViewModel(
         private val fingerprintRepository: FingerprintRepository
     ) : ChallengeGenerator {
 
-        private var mCallback: GenerateChallengeCallback? = null
-
-        override fun getCallback(): GenerateChallengeCallback? {
-            return mCallback
-        }
-
-        override fun setCallback(callback: GenerateChallengeCallback?) {
-            mCallback = callback
-        }
+        override var callback: GenerateChallengeCallback? = null
 
         override fun generateChallenge(userId: Int) {
-            val callback = mCallback
-            if (callback == null) {
+            callback?.let {
+                fingerprintRepository.generateChallenge(userId) {
+                        sensorId: Int, uid: Int, challenge: Long ->
+                    it.onChallengeGenerated(sensorId, uid, challenge)
+                }
+            } ?:run {
                 Log.e(TAG, "generateChallenge, null callback")
-                return
-            }
-
-            fingerprintRepository.generateChallenge(userId) {
-                sensorId: Int, uid: Int, challenge: Long ->
-                callback.onChallengeGenerated(
-                    sensorId,
-                    uid,
-                    challenge
-                )
             }
         }
 
@@ -160,7 +143,7 @@ class AutoCredentialViewModel(
         revokeGkPwHandle: Boolean,
         scope: CoroutineScope
     ) {
-        challengeGenerator.setCallback(object : GenerateChallengeCallback {
+        challengeGenerator.callback = object : GenerateChallengeCallback {
             override fun onChallengeGenerated(sensorId: Int, userId: Int, challenge: Long) {
                 var illegalStateExceptionCaught = false
                 try {
@@ -188,7 +171,7 @@ class AutoCredentialViewModel(
                     }
                 }
             }
-        })
+        }
         challengeGenerator.generateChallenge(userId)
     }
 
diff --git a/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModelTest.kt b/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModelTest.kt
index 3ae4951944b..976caed5863 100644
--- a/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModelTest.kt
+++ b/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/AutoCredentialViewModelTest.kt
@@ -507,20 +507,13 @@ class AutoCredentialViewModelTest {
         var challenge = CredentialModel.INVALID_CHALLENGE
         var callbackRunCount = 0
 
-        private var _callback: AutoCredentialViewModel.GenerateChallengeCallback? = null
-
-        override fun getCallback(): AutoCredentialViewModel.GenerateChallengeCallback? {
-            return _callback
-        }
-
-        override fun setCallback(callback: AutoCredentialViewModel.GenerateChallengeCallback?) {
-            _callback = callback
-        }
+        override var callback: AutoCredentialViewModel.GenerateChallengeCallback? = null
 
         override fun generateChallenge(userId: Int) {
-            val callback = _callback ?: return
-            callback.onChallengeGenerated(sensorId, this.userId, challenge)
-            ++callbackRunCount
+            callback?.let {
+                it.onChallengeGenerated(sensorId, this.userId, challenge)
+                ++callbackRunCount
+            }
         }
     }