diff --git a/biometrics/fingerprint/aidl/default/Fingerprint.cpp b/biometrics/fingerprint/aidl/default/Fingerprint.cpp
index 67dc34f..8028089 100644
--- a/biometrics/fingerprint/aidl/default/Fingerprint.cpp
+++ b/biometrics/fingerprint/aidl/default/Fingerprint.cpp
@@ -57,12 +57,10 @@
 ndk::ScopedAStatus Fingerprint::createSession(int32_t sensorId, int32_t userId,
                                               const std::shared_ptr<ISessionCallback>& cb,
                                               std::shared_ptr<ISession>* out) {
-    auto sessionSp = mSession.lock();
-    CHECK(sessionSp == nullptr || sessionSp->isClosed()) << "Open session already exists!";
+    CHECK(mSession == nullptr || mSession->isClosed()) << "Open session already exists!";
 
-    auto session = SharedRefBase::make<Session>(sensorId, userId, cb, mEngine.get(), &mWorker);
-    mSession = session;
-    *out = session;
+    mSession = SharedRefBase::make<Session>(sensorId, userId, cb, mEngine.get(), &mWorker);
+    *out = mSession;
     return ndk::ScopedAStatus::ok();
 }
 
diff --git a/biometrics/fingerprint/aidl/default/include/Fingerprint.h b/biometrics/fingerprint/aidl/default/include/Fingerprint.h
index 9b8eef8..9b43419 100644
--- a/biometrics/fingerprint/aidl/default/include/Fingerprint.h
+++ b/biometrics/fingerprint/aidl/default/include/Fingerprint.h
@@ -39,7 +39,7 @@
   private:
     std::unique_ptr<FakeFingerprintEngine> mEngine;
     WorkerThread mWorker;
-    std::weak_ptr<Session> mSession;
+    std::shared_ptr<Session> mSession;
 };
 
 }  // namespace aidl::android::hardware::biometrics::fingerprint
diff --git a/biometrics/fingerprint/aidl/vts/VtsHalBiometricsFingerprintTargetTest.cpp b/biometrics/fingerprint/aidl/vts/VtsHalBiometricsFingerprintTargetTest.cpp
index ddbc0fe..894fdfe 100644
--- a/biometrics/fingerprint/aidl/vts/VtsHalBiometricsFingerprintTargetTest.cpp
+++ b/biometrics/fingerprint/aidl/vts/VtsHalBiometricsFingerprintTargetTest.cpp
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #include <aidl/Gtest.h>
 #include <aidl/Vintf.h>
 #include <aidl/android/hardware/biometrics/fingerprint/BnFingerprint.h>
@@ -30,25 +31,35 @@
 constexpr int kUserId = 0;
 constexpr auto kCallbackTimeout = std::chrono::seconds(1);
 
-enum class SessionCallbackMethodName {
+enum class MethodName {
     kOnStateChanged,
 };
 
-struct SessionCallbackInvocation {
-    SessionCallbackMethodName method_name;
+struct Invocation {
+    MethodName methodName;
+    int32_t cookie;
     SessionState state;
 };
 
 class SessionCallback : public BnSessionCallback {
   public:
-    explicit SessionCallback(std::promise<SessionCallbackInvocation> invocation_promise)
-        : invocation_promise_(std::move(invocation_promise)) {}
+    explicit SessionCallback() : mIsPromiseValid(false) {}
 
-    ndk::ScopedAStatus onStateChanged(int32_t /*cookie*/, SessionState state) override {
-        SessionCallbackInvocation invocation = {};
-        invocation.method_name = SessionCallbackMethodName::kOnStateChanged;
+    void setPromise(std::promise<std::vector<Invocation>>&& promise) {
+        mPromise = std::move(promise);
+        mIsPromiseValid = true;
+    }
+
+    ndk::ScopedAStatus onStateChanged(int32_t cookie, SessionState state) override {
+        Invocation invocation = {};
+        invocation.methodName = MethodName::kOnStateChanged;
+        invocation.cookie = cookie;
         invocation.state = state;
-        invocation_promise_.set_value(invocation);
+        mInvocations.push_back(invocation);
+        if (state == SessionState::IDLING) {
+            assert(mIsPromiseValid);
+            mPromise.set_value(mInvocations);
+        }
         return ndk::ScopedAStatus::ok();
     }
 
@@ -73,26 +84,20 @@
         return ndk::ScopedAStatus::ok();
     }
 
-    ndk::ScopedAStatus onAuthenticationSucceeded(int32_t /*enrollmentId*/,
-                                       const keymaster::HardwareAuthToken& /*hat*/) override {
+    ndk::ScopedAStatus onAuthenticationSucceeded(
+            int32_t /*enrollmentId*/, const keymaster::HardwareAuthToken& /*hat*/) override {
         return ndk::ScopedAStatus::ok();
     }
 
-    ndk::ScopedAStatus onAuthenticationFailed() override {
-        return ndk::ScopedAStatus::ok();
-    }
+    ndk::ScopedAStatus onAuthenticationFailed() override { return ndk::ScopedAStatus::ok(); }
 
     ndk::ScopedAStatus onLockoutTimed(int64_t /*durationMillis*/) override {
         return ndk::ScopedAStatus::ok();
     }
 
-    ndk::ScopedAStatus onLockoutPermanent() override {
-        return ndk::ScopedAStatus::ok();
-    }
+    ndk::ScopedAStatus onLockoutPermanent() override { return ndk::ScopedAStatus::ok(); }
 
-    ndk::ScopedAStatus onLockoutCleared() override {
-        return ndk::ScopedAStatus::ok();
-    }
+    ndk::ScopedAStatus onLockoutCleared() override { return ndk::ScopedAStatus::ok(); }
 
     ndk::ScopedAStatus onInteractionDetected() override { return ndk::ScopedAStatus::ok(); }
 
@@ -115,7 +120,9 @@
     }
 
   private:
-    std::promise<SessionCallbackInvocation> invocation_promise_;
+    bool mIsPromiseValid;
+    std::vector<Invocation> mInvocations;
+    std::promise<std::vector<Invocation>> mPromise;
 };
 
 class Fingerprint : public testing::TestWithParam<std::string> {
@@ -123,28 +130,40 @@
     void SetUp() override {
         AIBinder* binder = AServiceManager_waitForService(GetParam().c_str());
         ASSERT_NE(binder, nullptr);
-        hal_ = IFingerprint::fromBinder(ndk::SpAIBinder(binder));
+        mHal = IFingerprint::fromBinder(ndk::SpAIBinder(binder));
     }
 
-    std::shared_ptr<IFingerprint> hal_;
+    std::shared_ptr<IFingerprint> mHal;
 };
 
 TEST_P(Fingerprint, AuthenticateTest) {
-    std::promise<SessionCallbackInvocation> invocation_promise;
-    std::future<SessionCallbackInvocation> invocation_future = invocation_promise.get_future();
-    std::shared_ptr<SessionCallback> session_cb =
-            ndk::SharedRefBase::make<SessionCallback>(std::move(invocation_promise));
+    // Prepare the callback
+    std::promise<std::vector<Invocation>> promise;
+    auto future = promise.get_future();
+    std::shared_ptr<SessionCallback> cb = ndk::SharedRefBase::make<SessionCallback>();
+    cb->setPromise(std::move(promise));
 
+    // Create a session
     std::shared_ptr<ISession> session;
-    ASSERT_TRUE(hal_->createSession(kSensorId, kUserId, session_cb, &session).isOk());
+    ASSERT_TRUE(mHal->createSession(kSensorId, kUserId, cb, &session).isOk());
 
-    std::shared_ptr<common::ICancellationSignal> cancel_cb;
-    ASSERT_TRUE(session->authenticate(0, 0, &cancel_cb).isOk());
-    ASSERT_EQ(invocation_future.wait_for(kCallbackTimeout), std::future_status::ready);
+    // Call authenticate
+    int32_t cookie = 123;
+    std::shared_ptr<common::ICancellationSignal> cancellationSignal;
+    ASSERT_TRUE(session->authenticate(cookie, 0, &cancellationSignal).isOk());
 
-    SessionCallbackInvocation invocation = invocation_future.get();
-    EXPECT_EQ(invocation.method_name, SessionCallbackMethodName::kOnStateChanged);
-    EXPECT_EQ(invocation.state, SessionState::AUTHENTICATING);
+    // Get the results
+    ASSERT_TRUE(future.wait_for(kCallbackTimeout) == std::future_status::ready);
+    std::vector<Invocation> invocations = future.get();
+
+    // Close the session
+    ASSERT_TRUE(session->close(0).isOk());
+
+    ASSERT_FALSE(invocations.empty());
+    EXPECT_EQ(invocations.front().methodName, MethodName::kOnStateChanged);
+    EXPECT_EQ(invocations.front().state, SessionState::AUTHENTICATING);
+    EXPECT_EQ(invocations.back().methodName, MethodName::kOnStateChanged);
+    EXPECT_EQ(invocations.back().state, SessionState::IDLING);
 }
 
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(Fingerprint);
@@ -154,6 +173,7 @@
         ::android::PrintInstanceNameToString);
 
 }  // namespace
+}  // namespace aidl::android::hardware::biometrics::fingerprint
 
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
@@ -161,5 +181,3 @@
     ABinderProcess_startThreadPool();
     return RUN_ALL_TESTS();
 }
-
-}  // namespace aidl::android::hardware::biometrics::fingerprint