diff --git a/tests/StagedInstallTest/Android.bp b/tests/StagedInstallTest/Android.bp
index 2e57467e55fec..c679d04876297 100644
--- a/tests/StagedInstallTest/Android.bp
+++ b/tests/StagedInstallTest/Android.bp
@@ -25,17 +25,24 @@ android_test_helper_app {
     name: "StagedInstallInternalTestApp",
     manifest: "app/AndroidManifest.xml",
     srcs: ["app/src/**/*.java"],
-    static_libs: ["androidx.test.rules", "cts-install-lib"],
+    static_libs: [
+        "androidx.test.rules",
+        "cts-install-lib",
+    ],
     test_suites: ["general-tests"],
     java_resources: [
         ":com.android.apex.apkrollback.test_v2",
+        ":StagedInstallTestApexV2_WrongSha",
     ],
 }
 
 java_test_host {
     name: "StagedInstallInternalTest",
     srcs: ["src/**/*.java"],
-    libs: ["tradefed", "cts-shim-host-lib"],
+    libs: [
+        "tradefed",
+        "cts-shim-host-lib",
+    ],
     static_libs: [
         "testng",
         "compatibility-tradefed",
diff --git a/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java b/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
index ad8aac17d8443..e633c87d7fbbb 100644
--- a/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
@@ -17,6 +17,7 @@
 package com.android.tests.stagedinstallinternal;
 
 import static com.android.cts.install.lib.InstallUtils.getPackageInstaller;
+import static com.android.cts.shim.lib.ShimPackage.SHIM_APEX_PACKAGE_NAME;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -50,6 +51,9 @@ public class StagedInstallInternalTest {
     private static final String APK_IN_APEX_TESTAPEX_NAME = "com.android.apex.apkrollback.test";
     private static final TestApp TEST_APEX_WITH_APK_V2 = new TestApp("TestApexWithApkV2",
             APK_IN_APEX_TESTAPEX_NAME, 2, /*isApex*/true, APK_IN_APEX_TESTAPEX_NAME + "_v2.apex");
+    private static final TestApp APEX_WRONG_SHA_V2 = new TestApp(
+            "ApexWrongSha2", SHIM_APEX_PACKAGE_NAME, 2, /*isApex*/true,
+            "com.android.apex.cts.shim.v2_wrong_sha.apex");
 
     private File mTestStateFile = new File(
             InstrumentationRegistry.getInstrumentation().getContext().getFilesDir(),
@@ -127,6 +131,26 @@ public class StagedInstallInternalTest {
         InstallUtils.getPackageInstaller().abandonSession(id4);
     }
 
+    @Test
+    public void testStagedSessionShouldCleanUpOnVerificationFailure() throws Exception {
+        InstallUtils.commitExpectingFailure(AssertionError.class, "apexd verification failed",
+                Install.single(APEX_WRONG_SHA_V2).setStaged());
+    }
+
+    @Test
+    public void testStagedSessionShouldCleanUpOnOnSuccess_Commit() throws Exception {
+        int sessionId = Install.single(TestApp.A1).setStaged().commit();
+        storeSessionId(sessionId);
+    }
+
+    @Test
+    public void testStagedSessionShouldCleanUpOnOnSuccess_Verify() throws Exception {
+        int sessionId = retrieveLastSessionId();
+        PackageInstaller.SessionInfo info = InstallUtils.getStagedSessionInfo(sessionId);
+        assertThat(info).isNotNull();
+        assertThat(info.isStagedSessionApplied()).isTrue();
+    }
+
     @Test
     public void testStagedInstallationShouldCleanUpOnValidationFailure() throws Exception {
         InstallUtils.commitExpectingFailure(AssertionError.class, "INSTALL_FAILED_INVALID_APK",
diff --git a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
index 8dc53ac267152..9fd190cc12fe6 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -44,7 +44,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.io.File;
-import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -251,6 +250,28 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         assertThat(after).isEqualTo(before);
     }
 
+    @Test
+    public void testStagedSessionShouldCleanUpOnVerificationFailure() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mHostUtils.isApexUpdateSupported());
+        List<String> before = getStagingDirectories();
+        runPhase("testStagedSessionShouldCleanUpOnVerificationFailure");
+        List<String> after = getStagingDirectories();
+        assertThat(after).isEqualTo(before);
+    }
+
+    @Test
+    @LargeTest
+    public void testStagedSessionShouldCleanUpOnOnSuccess() throws Exception {
+        List<String> before = getStagingDirectories();
+        runPhase("testStagedSessionShouldCleanUpOnOnSuccess_Commit");
+        assertThat(getStagingDirectories()).isNotEqualTo(before);
+        getDevice().reboot();
+        runPhase("testStagedSessionShouldCleanUpOnOnSuccess_Verify");
+        List<String> after = getStagingDirectories();
+        assertThat(after).isEqualTo(before);
+    }
+
     @Test
     public void testStagedInstallationShouldCleanUpOnValidationFailure() throws Exception {
         List<String> before = getStagingDirectories();
@@ -273,12 +294,14 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         //create random directories in /data/app-staging folder
         getDevice().enableAdbRoot();
         getDevice().executeShellCommand("mkdir /data/app-staging/session_123");
-        getDevice().executeShellCommand("mkdir /data/app-staging/random_name");
+        getDevice().executeShellCommand("mkdir /data/app-staging/session_456");
         getDevice().disableAdbRoot();
 
-        assertThat(getStagingDirectories()).isNotEmpty();
+        assertThat(getStagingDirectories()).contains("session_123");
+        assertThat(getStagingDirectories()).contains("session_456");
         getDevice().reboot();
-        assertThat(getStagingDirectories()).isEmpty();
+        assertThat(getStagingDirectories()).doesNotContain("session_123");
+        assertThat(getStagingDirectories()).doesNotContain("session_456");
     }
 
     @Test
@@ -304,9 +327,6 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
                     .stream().filter(entry -> entry.getName().matches("session_\\d+"))
                     .map(entry -> entry.getName())
                     .collect(Collectors.toList());
-        } catch (Exception e) {
-            // Return an empty list if any error
-            return Collections.EMPTY_LIST;
         } finally {
             getDevice().disableAdbRoot();
         }