Merge changes from topic "compos-cherry-pick" into stage-aosp-master

* changes:
  Hook the new APIs in StagingManager to PackageManagerNative service
  Notify StagedApexObservers when there is a change in set of staged APEX
  Open up new API in StagingManager to get information about staged APEX
  Connect the new API from ApexService to ApexManager

services/core/java/com/android/server/pm/ApexManager.java

@@ -248,6 +248,14 @@ public abstract class ApexManager {
     abstract ApexInfoList submitStagedSession(ApexSessionParams params)
             throws PackageManagerException;
 
+    /**
+     * Returns {@code ApeInfo} about apex sessions that have been marked ready via
+     * {@link #markStagedSessionReady(int)}
+     *
+     * Returns empty array if there is no staged apex session or if there is any error.
+     */
+    abstract ApexInfo[] getStagedApexInfos(ApexSessionParams params);
+
     /**
      * Mark a staged session previously submitted using {@code submitStagedSession} as ready to be
      * applied at next reboot.
@@ -756,6 +764,19 @@ public abstract class ApexManager {
             }
         }
 
+        @Override
+        ApexInfo[] getStagedApexInfos(ApexSessionParams params) {
+            try {
+                return waitForApexService().getStagedApexInfos(params);
+            } catch (RemoteException re) {
+                Slog.w(TAG, "Unable to contact apexservice" + re.getMessage());
+                throw new RuntimeException(re);
+            } catch (Exception e) {
+                Slog.w(TAG, "Failed to collect staged apex infos" + e.getMessage());
+                return new ApexInfo[0];
+            }
+        }
+
         @Override
         void markStagedSessionReady(int sessionId) throws PackageManagerException {
             try {
@@ -1246,6 +1267,11 @@ public abstract class ApexManager {
                     "Device doesn't support updating APEX");
         }
 
+        @Override
+        ApexInfo[] getStagedApexInfos(ApexSessionParams params) {
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         void markStagedSessionReady(int sessionId) {
             throw new UnsupportedOperationException();

services/core/java/com/android/server/pm/PackageInstallerService.java

@@ -260,6 +260,10 @@ public class PackageInstallerService extends IPackageInstaller.Stub implements
                 new Lifecycle(context, this));
     }
 
+    StagingManager getStagingManager() {
+        return mStagingManager;
+    }
+
     boolean okToSendBroadcasts()  {
         return mOkToSendBroadcasts;
     }

services/core/java/com/android/server/pm/PackageManagerService.java

@@ -186,6 +186,7 @@ import android.content.pm.IPackageManager;
 import android.content.pm.IPackageManagerNative;
 import android.content.pm.IPackageMoveObserver;
 import android.content.pm.IPackageStatsObserver;
+import android.content.pm.IStagedApexObserver;
 import android.content.pm.IncrementalStatesInfo;
 import android.content.pm.InstallSourceInfo;
 import android.content.pm.InstantAppInfo;
@@ -227,6 +228,7 @@ import android.content.pm.ServiceInfo;
 import android.content.pm.SharedLibraryInfo;
 import android.content.pm.Signature;
 import android.content.pm.SigningInfo;
+import android.content.pm.StagedApexInfo;
 import android.content.pm.SuspendDialogInfo;
 import android.content.pm.TestUtilityService;
 import android.content.pm.UserInfo;
@@ -27003,6 +27005,29 @@ public class PackageManagerService extends IPackageManager.Stub
         public boolean hasSystemFeature(String featureName, int version) {
             return PackageManagerService.this.hasSystemFeature(featureName, version);
         }
+
+        @Override
+        public void registerStagedApexObserver(IStagedApexObserver observer) {
+            mInstallerService.getStagingManager().registerStagedApexObserver(observer);
+        }
+
+        @Override
+        public void unregisterStagedApexObserver(IStagedApexObserver observer) {
+            mInstallerService.getStagingManager().unregisterStagedApexObserver(observer);
+        }
+
+        @Override
+        public String[] getStagedApexModuleNames() {
+            return mInstallerService.getStagingManager()
+                    .getStagedApexModuleNames().toArray(new String[0]);
+        }
+
+        @Override
+        @Nullable
+        public StagedApexInfo getStagedApexInfo(String moduleName) {
+            return mInstallerService.getStagingManager().getStagedApexInfo(moduleName);
+        }
+
     }
 
     private AndroidPackage getPackage(String packageName) {

services/core/java/com/android/server/pm/StagingManager.java

@@ -17,6 +17,7 @@
 package com.android.server.pm;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.apex.ApexInfo;
 import android.apex.ApexInfoList;
 import android.apex.ApexSessionInfo;
@@ -28,7 +29,9 @@ import android.content.IIntentSender;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.IntentSender;
+import android.content.pm.ApexStagedEvent;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.IStagedApexObserver;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
@@ -38,6 +41,7 @@ import android.content.pm.PackageManagerInternal;
 import android.content.pm.PackageParser.PackageParserException;
 import android.content.pm.PackageParser.SigningDetails;
 import android.content.pm.PackageParser.SigningDetails.SignatureSchemeVersion;
+import android.content.pm.StagedApexInfo;
 import android.content.pm.parsing.PackageInfoWithoutStateUtils;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
@@ -52,6 +56,7 @@ import android.os.SystemProperties;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.text.TextUtils;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Slog;
@@ -80,7 +85,9 @@ import java.io.File;
 import java.io.FileReader;
 import java.io.FileWriter;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -98,7 +105,8 @@ public class StagingManager {
     private final ApexManager mApexManager;
     private final PowerManager mPowerManager;
     private final Context mContext;
-    private final PreRebootVerificationHandler mPreRebootVerificationHandler;
+    @VisibleForTesting
+    final PreRebootVerificationHandler mPreRebootVerificationHandler;
     private final Supplier<PackageParser2> mPackageParserSupplier;
 
     private final File mFailureReasonFile = new File("/metadata/staged-install/failure_reason.txt");
@@ -114,6 +122,9 @@ public class StagingManager {
     @GuardedBy("mSuccessfulStagedSessionIds")
     private final List<Integer> mSuccessfulStagedSessionIds = new ArrayList<>();
 
+    @GuardedBy("mStagedApexObservers")
+    private final List<IStagedApexObserver> mStagedApexObservers = new ArrayList<>();
+
     interface StagedSession {
         boolean isMultiPackage();
         boolean isApexSession();
@@ -198,6 +209,35 @@ public class StagingManager {
         mApexManager.markBootCompleted();
     }
 
+    void registerStagedApexObserver(IStagedApexObserver observer) {
+        if (observer == null) {
+            return;
+        }
+        if  (observer.asBinder() != null) {
+            try {
+                observer.asBinder().linkToDeath(new IBinder.DeathRecipient() {
+                    @Override
+                    public void binderDied() {
+                        synchronized (mStagedApexObservers) {
+                            mStagedApexObservers.remove(observer);
+                        }
+                    }
+                }, 0);
+            } catch (RemoteException re) {
+                Slog.w(TAG, re.getMessage());
+            }
+        }
+        synchronized (mStagedApexObservers) {
+            mStagedApexObservers.add(observer);
+        }
+    }
+
+    void unregisterStagedApexObserver(IStagedApexObserver observer) {
+        synchronized (mStagedApexObservers) {
+            mStagedApexObservers.remove(observer);
+        }
+    }
+
     /**
      * Validates the signature used to sign the container of the new apex package
      *
@@ -836,6 +876,9 @@ public class StagingManager {
                 // Also, cleaning up the stageDir prevents the apex from being activated.
                 Slog.e(TAG, "Failed to abort apex session " + session.sessionId());
             }
+            if (session.containsApexSession()) {
+                notifyStagedApexObservers();
+            }
         }
 
         // Session was successfully aborted from apexd (if required) and pre-reboot verification
@@ -1134,7 +1177,107 @@ public class StagingManager {
         mPreRebootVerificationHandler.notifyPreRebootVerification_Apk_Complete(session);
     }
 
-    private final class PreRebootVerificationHandler extends Handler {
+    /**
+     * Returns ApexInfo about APEX contained inside the session as a {@code Map<String, ApexInfo>},
+     * where the key of the map is the module name of the ApexInfo.
+     *
+     * Returns an empty map if there is any error.
+     */
+    @VisibleForTesting
+    @NonNull
+    Map<String, ApexInfo> getStagedApexInfos(@NonNull StagedSession session) {
+        Preconditions.checkArgument(session != null, "Session is null");
+        Preconditions.checkArgument(!session.hasParentSessionId(),
+                session.sessionId() + " session has parent session");
+        Preconditions.checkArgument(session.containsApexSession(),
+                session.sessionId() + " session does not contain apex");
+
+        // Even if caller calls this method on ready session, the session could be abandoned
+        // right after this method is called.
+        if (!session.isSessionReady() || session.isDestroyed()) {
+            return Collections.emptyMap();
+        }
+
+        ApexSessionParams params = new ApexSessionParams();
+        params.sessionId = session.sessionId();
+        final IntArray childSessionIds = new IntArray();
+        if (session.isMultiPackage()) {
+            for (StagedSession s : session.getChildSessions()) {
+                if (s.isApexSession()) {
+                    childSessionIds.add(s.sessionId());
+                }
+            }
+        }
+        params.childSessionIds = childSessionIds.toArray();
+
+        ApexInfo[] infos = mApexManager.getStagedApexInfos(params);
+        Map<String, ApexInfo> result = new ArrayMap<>();
+        for (ApexInfo info : infos) {
+            result.put(info.moduleName, info);
+        }
+        return result;
+    }
+
+    /**
+     * Returns apex module names of all packages that are staged ready
+     */
+    List<String> getStagedApexModuleNames() {
+        List<String> result = new ArrayList<>();
+        synchronized (mStagedSessions) {
+            for (int i = 0; i < mStagedSessions.size(); i++) {
+                final StagedSession session = mStagedSessions.valueAt(i);
+                if (!session.isSessionReady() || session.isDestroyed()
+                        || session.hasParentSessionId() || !session.containsApexSession()) {
+                    continue;
+                }
+                result.addAll(getStagedApexInfos(session).keySet());
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Returns ApexInfo of the {@code moduleInfo} provided if it is staged, otherwise returns null.
+     */
+    @Nullable
+    StagedApexInfo getStagedApexInfo(String moduleName) {
+        synchronized (mStagedSessions) {
+            for (int i = 0; i < mStagedSessions.size(); i++) {
+                final StagedSession session = mStagedSessions.valueAt(i);
+                if (!session.isSessionReady() || session.isDestroyed()
+                        || session.hasParentSessionId() || !session.containsApexSession()) {
+                    continue;
+                }
+                ApexInfo ai = getStagedApexInfos(session).get(moduleName);
+                if (ai != null) {
+                    StagedApexInfo info = new StagedApexInfo();
+                    info.moduleName = ai.moduleName;
+                    info.diskImagePath = ai.modulePath;
+                    info.versionCode = ai.versionCode;
+                    info.versionName = ai.versionName;
+                    return info;
+                }
+            }
+        }
+        return null;
+    }
+
+    private void notifyStagedApexObservers() {
+        synchronized (mStagedApexObservers) {
+            for (IStagedApexObserver observer : mStagedApexObservers) {
+                ApexStagedEvent event = new ApexStagedEvent();
+                event.stagedApexModuleNames = getStagedApexModuleNames().toArray(new String[0]);
+                try {
+                    observer.onApexStaged(event);
+                } catch (RemoteException re) {
+                    Slog.w(TAG, "Failed to contact the observer " + re.getMessage());
+                }
+            }
+        }
+    }
+
+    @VisibleForTesting
+    final class PreRebootVerificationHandler extends Handler {
         // Hold sessions before handler gets ready to do the verification.
         private List<StagedSession> mPendingSessions;
         private boolean mIsReady;
@@ -1160,7 +1303,8 @@ public class StagingManager {
         private static final int MSG_PRE_REBOOT_VERIFICATION_START = 1;
         private static final int MSG_PRE_REBOOT_VERIFICATION_APEX = 2;
         private static final int MSG_PRE_REBOOT_VERIFICATION_APK = 3;
-        private static final int MSG_PRE_REBOOT_VERIFICATION_END = 4;
+        @VisibleForTesting
+        static final int MSG_PRE_REBOOT_VERIFICATION_END = 4;
 
         @Override
         public void handleMessage(Message msg) {
@@ -1402,6 +1546,7 @@ public class StagingManager {
                 if (hasApex) {
                     try {
                         mApexManager.markStagedSessionReady(session.sessionId());
+                        notifyStagedApexObservers();
                     } catch (PackageManagerException e) {
                         session.setSessionFailed(e.error, e.getMessage());
                         return;

services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java

@@ -18,6 +18,7 @@ package com.android.server.pm;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -28,17 +29,23 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertThrows;
 
+import android.apex.ApexInfo;
 import android.apex.ApexSessionInfo;
+import android.apex.ApexSessionParams;
 import android.content.Context;
 import android.content.IntentSender;
+import android.content.pm.ApexStagedEvent;
+import android.content.pm.IStagedApexObserver;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
 import android.content.pm.PackageInstaller.SessionInfo.StagedSessionErrorCode;
+import android.content.pm.StagedApexInfo;
+import android.os.Message;
 import android.os.SystemProperties;
 import android.os.storage.IStorageManager;
 import android.platform.test.annotations.Presubmit;
+import android.util.IntArray;
 import android.util.SparseArray;
 
 import com.android.dx.mockito.inline.extended.ExtendedMockito;
@@ -53,15 +60,20 @@ import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.mockito.MockitoSession;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.quality.Strictness;
+import org.mockito.stubbing.Answer;
 
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Predicate;
 
 @Presubmit
@@ -519,6 +531,281 @@ public class StagingManagerTest {
         assertThat(mStagingManager.getSessionIdByPackageName("com.bar")).isEqualTo(-1);
     }
 
+    @Test
+    public void getStagedApexInfos_validatePreConditions() throws Exception {
+        // Invalid session: null session
+        {
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(null));
+            assertThat(thrown).hasMessageThat().contains("Session is null");
+        }
+        // Invalid session: has parent
+        {
+            FakeStagedSession session = new FakeStagedSession(241);
+            session.setParentSessionId(239);
+            session.setSessionReady();
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(session));
+            assertThat(thrown).hasMessageThat().contains("241 session has parent");
+        }
+
+        // Invalid session: does not contain apex
+        {
+            FakeStagedSession session = new FakeStagedSession(241);
+            session.setSessionReady();
+            // Call and verify
+            IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class,
+                    () -> mStagingManager.getStagedApexInfos(session));
+            assertThat(thrown).hasMessageThat().contains("241 session does not contain apex");
+        }
+        // Invalid session: not ready
+        {
+            FakeStagedSession session = new FakeStagedSession(239);
+            session.setIsApex(true);
+            // Call and verify
+            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            assertThat(result).isEmpty();
+        }
+        // Invalid session: destroyed
+        {
+            FakeStagedSession session = new FakeStagedSession(240);
+            session.setSessionReady();
+            session.setIsApex(true);
+            session.setDestroyed(true);
+            // Call and verify
+            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            assertThat(result).isEmpty();
+        }
+    }
+
+    private ApexInfo[] getFakeApexInfo(List<String> moduleNames) {
+        List<ApexInfo> result = new ArrayList<>();
+        for (String moduleName : moduleNames) {
+            ApexInfo info = new ApexInfo();
+            info.moduleName = moduleName;
+            result.add(info);
+        }
+        return result.toArray(new ApexInfo[0]);
+    }
+
+    @Test
+    public void getStagedApexInfos_nonParentSession() throws Exception {
+        FakeStagedSession validSession = new FakeStagedSession(239);
+        validSession.setIsApex(true);
+        validSession.setSessionReady();
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Call and verify
+        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(validSession);
+        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0]);
+
+        ArgumentCaptor<ApexSessionParams> argumentCaptor =
+                ArgumentCaptor.forClass(ApexSessionParams.class);
+        verify(mApexManager, times(1)).getStagedApexInfos(argumentCaptor.capture());
+        ApexSessionParams params = argumentCaptor.getValue();
+        assertThat(params.sessionId).isEqualTo(239);
+    }
+
+    @Test
+    public void getStagedApexInfos_parentSession() throws Exception {
+        FakeStagedSession childSession1 = new FakeStagedSession(201);
+        childSession1.setIsApex(true);
+        FakeStagedSession childSession2 = new FakeStagedSession(202);
+        childSession2.setIsApex(true);
+        FakeStagedSession nonApexChild = new FakeStagedSession(203);
+        FakeStagedSession parentSession = new FakeStagedSession(239,
+                Arrays.asList(childSession1, childSession2, nonApexChild));
+        parentSession.setSessionReady();
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1", "module2"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Call and verify
+        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(parentSession);
+        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0],
+                fakeApexInfos[1].moduleName, fakeApexInfos[1]);
+
+        ArgumentCaptor<ApexSessionParams> argumentCaptor =
+                ArgumentCaptor.forClass(ApexSessionParams.class);
+        verify(mApexManager, times(1)).getStagedApexInfos(argumentCaptor.capture());
+        ApexSessionParams params = argumentCaptor.getValue();
+        assertThat(params.sessionId).isEqualTo(239);
+        assertThat(params.childSessionIds).asList().containsExactly(201, 202);
+    }
+
+    @Test
+    public void getStagedApexModuleNames_returnsStagedApexModules() throws Exception {
+        FakeStagedSession validSession1 = new FakeStagedSession(239);
+        validSession1.setIsApex(true);
+        validSession1.setSessionReady();
+        mStagingManager.createSession(validSession1);
+
+        FakeStagedSession childSession1 = new FakeStagedSession(123);
+        childSession1.setIsApex(true);
+        FakeStagedSession childSession2 = new FakeStagedSession(124);
+        childSession2.setIsApex(true);
+        FakeStagedSession nonApexChild = new FakeStagedSession(125);
+        FakeStagedSession parentSession = new FakeStagedSession(240,
+                Arrays.asList(childSession1, childSession2, nonApexChild));
+        parentSession.setSessionReady();
+        mStagingManager.createSession(parentSession);
+
+        mockApexManagerGetStagedApexInfoWithSessionId();
+
+        List<String> result = mStagingManager.getStagedApexModuleNames();
+        assertThat(result).containsExactly("239", "123", "124");
+        verify(mApexManager, times(2)).getStagedApexInfos(any());
+    }
+
+    // Make mApexManager return ApexInfo with same module name as the sessionId
+    // of the parameter that was passed into it
+    private void mockApexManagerGetStagedApexInfoWithSessionId() {
+        when(mApexManager.getStagedApexInfos(any())).thenAnswer(new Answer<ApexInfo[]>() {
+            @Override
+            public ApexInfo[] answer(InvocationOnMock invocation) throws Throwable {
+                Object[] args = invocation.getArguments();
+                ApexSessionParams params = (ApexSessionParams) args[0];
+                IntArray sessionsToProcess = new IntArray();
+                if (params.childSessionIds.length == 0) {
+                    sessionsToProcess.add(params.sessionId);
+                } else {
+                    sessionsToProcess.addAll(params.childSessionIds);
+                }
+                List<ApexInfo> result = new ArrayList<>();
+                for (int session : sessionsToProcess.toArray()) {
+                    ApexInfo info = new ApexInfo();
+                    info.moduleName = String.valueOf(session);
+                    result.add(info);
+                }
+                return result.toArray(new ApexInfo[0]);
+            }
+        });
+    }
+
+    @Test
+    public void getStagedApexInfo() throws Exception {
+        FakeStagedSession validSession1 = new FakeStagedSession(239);
+        validSession1.setIsApex(true);
+        validSession1.setSessionReady();
+        mStagingManager.createSession(validSession1);
+        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1"));
+        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
+
+        // Verify null is returned if module name is not found
+        StagedApexInfo result = mStagingManager.getStagedApexInfo("not found");
+        assertThat(result).isNull();
+        verify(mApexManager, times(1)).getStagedApexInfos(any());
+        // Otherwise, the correct object is returned
+        result = mStagingManager.getStagedApexInfo("module1");
+        assertThat(result.moduleName).isEqualTo(fakeApexInfos[0].moduleName);
+        assertThat(result.diskImagePath).isEqualTo(fakeApexInfos[0].modulePath);
+        assertThat(result.versionCode).isEqualTo(fakeApexInfos[0].versionCode);
+        assertThat(result.versionName).isEqualTo(fakeApexInfos[0].versionName);
+        verify(mApexManager, times(2)).getStagedApexInfos(any());
+    }
+
+    @Test
+    public void registeredStagedApexObserverIsNotifiedOnPreRebootVerificationCompletion()
+            throws Exception {
+        // Register observer
+        IStagedApexObserver observer = Mockito.mock(IStagedApexObserver.class);
+        mStagingManager.registerStagedApexObserver(observer);
+
+        // Create one staged session and trigger end of pre-reboot verification
+        {
+            FakeStagedSession session = new FakeStagedSession(239);
+            session.setIsApex(true);
+            mStagingManager.createSession(session);
+
+            mockApexManagerGetStagedApexInfoWithSessionId();
+            triggerEndOfPreRebootVerification(session);
+
+            assertThat(session.isSessionReady()).isTrue();
+            ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
+                    ApexStagedEvent.class);
+            verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
+            assertThat(argumentCaptor.getValue().stagedApexModuleNames).isEqualTo(
+                    new String[]{"239"});
+        }
+
+        // Create another staged session and verify observers are notified of union
+        {
+            Mockito.clearInvocations(observer);
+            FakeStagedSession session = new FakeStagedSession(240);
+            session.setIsApex(true);
+            mStagingManager.createSession(session);
+
+            triggerEndOfPreRebootVerification(session);
+
+            assertThat(session.isSessionReady()).isTrue();
+            ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
+                    ApexStagedEvent.class);
+            verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
+            assertThat(argumentCaptor.getValue().stagedApexModuleNames).isEqualTo(
+                    new String[]{"239", "240"});
+        }
+
+        // Finally, verify that once unregistered, observer is not notified
+        mStagingManager.unregisterStagedApexObserver(observer);
+        {
+            Mockito.clearInvocations(observer);
+            FakeStagedSession session = new FakeStagedSession(241);
+            session.setIsApex(true);
+            mStagingManager.createSession(session);
+
+            triggerEndOfPreRebootVerification(session);
+
+            assertThat(session.isSessionReady()).isTrue();
+            verify(observer, never()).onApexStaged(any());
+        }
+    }
+
+    @Test
+    public void registeredStagedApexObserverIsNotifiedOnSessionAbandon() throws Exception {
+        // Register observer
+        IStagedApexObserver observer = Mockito.mock(IStagedApexObserver.class);
+        mStagingManager.registerStagedApexObserver(observer);
+
+        // Create a ready session and abandon it
+        FakeStagedSession session = new FakeStagedSession(239);
+        session.setIsApex(true);
+        session.setSessionReady();
+        session.setDestroyed(true);
+        mStagingManager.createSession(session);
+
+        mStagingManager.abortCommittedSession(session);
+
+        assertThat(session.isSessionReady()).isTrue();
+        ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
+                ApexStagedEvent.class);
+        verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
+        assertThat(argumentCaptor.getValue().stagedApexModuleNames).hasLength(0);
+    }
+
+    @Test
+    public void stagedApexObserverIsOnlyCalledForApexSessions() throws Exception {
+        IStagedApexObserver observer = Mockito.mock(IStagedApexObserver.class);
+        mStagingManager.registerStagedApexObserver(observer);
+
+        //  Trigger end of pre-reboot verification
+        FakeStagedSession session = new FakeStagedSession(239);
+        mStagingManager.createSession(session);
+
+        triggerEndOfPreRebootVerification(session);
+        assertThat(session.isSessionReady()).isTrue();
+        verify(observer, never()).onApexStaged(any());
+    }
+
+    private void triggerEndOfPreRebootVerification(StagingManager.StagedSession session) {
+        StagingManager.PreRebootVerificationHandler handler =
+                mStagingManager.mPreRebootVerificationHandler;
+        Message msg =  handler.obtainMessage(
+                handler.MSG_PRE_REBOOT_VERIFICATION_END, session.sessionId(), -1, session);
+        handler.handleMessage(msg);
+    }
+
     private StagingManager.StagedSession createSession(int sessionId, String packageName,
             long committedMillis) {
         PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
@@ -582,10 +869,16 @@ public class StagingManagerTest {
         private String mPackageName;
         private boolean mIsAbandonded = false;
         private boolean mPreRebootVerificationStarted = false;
-        private final List<StagingManager.StagedSession> mChildSessions = new ArrayList<>();
+        private final List<StagingManager.StagedSession> mChildSessions;
 
         private FakeStagedSession(int sessionId) {
             mSessionId = sessionId;
+            mChildSessions = new ArrayList<>();
+        }
+
+        private FakeStagedSession(int sessionId, List<StagingManager.StagedSession> childSessions) {
+            mSessionId = sessionId;
+            mChildSessions = childSessions;
         }
 
         private void setParentSessionId(int parentSessionId) {
@@ -775,9 +1068,7 @@ public class StagingManagerTest {
         }
 
         @Override
-        public void notifyEndPreRebootVerification() {
+        public void notifyEndPreRebootVerification() {}
-            throw new UnsupportedOperationException();
-        }
 
         @Override
         public void verifySession() {

services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java

@@ -40,6 +40,7 @@ import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.os.RemoteException;
+import android.os.ServiceSpecificException;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.filters.SmallTest;
@@ -231,6 +232,21 @@ public class ApexManagerTest {
                 () -> mApexManager.submitStagedSession(testParamsWithChildren()));
     }
 
+    @Test
+    public void testGetStagedApexInfos_throwRunTimeException() throws RemoteException {
+        doThrow(RemoteException.class).when(mApexService).getStagedApexInfos(any());
+
+        assertThrows(RuntimeException.class,
+                () -> mApexManager.getStagedApexInfos(testParamsWithChildren()));
+    }
+
+    @Test
+    public void testGetStagedApexInfos_returnsEmptyArrayOnError() throws RemoteException {
+        doThrow(ServiceSpecificException.class).when(mApexService).getStagedApexInfos(any());
+
+        assertThat(mApexManager.getStagedApexInfos(testParamsWithChildren())).hasLength(0);
+    }
+
     @Test
     public void testMarkStagedSessionReady_throwPackageManagerException() throws RemoteException {
         doAnswer(invocation -> {

tests/StagedInstallTest/Android.bp

@@ -37,6 +37,7 @@ android_test_helper_app {
         ":test.rebootless_apex_v1",
         ":test.rebootless_apex_v2",
     ],
+    platform_apis: true,
 }
 
 java_test_host {

tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java

@@ -17,16 +17,27 @@
 package com.android.tests.stagedinstallinternal;
 
 import static com.android.cts.install.lib.InstallUtils.getPackageInstaller;
+import static com.android.cts.install.lib.InstallUtils.waitForSessionReady;
 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;
 
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+
 import android.Manifest;
+import android.content.pm.ApexStagedEvent;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.IPackageManagerNative;
+import android.content.pm.IStagedApexObserver;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageManager;
+import android.content.pm.StagedApexInfo;
+import android.os.IBinder;
+import android.os.ServiceManager;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 
@@ -39,6 +50,8 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
 
 import java.io.BufferedReader;
 import java.io.BufferedWriter;
@@ -401,9 +414,73 @@ public class StagedInstallInternalTest {
                 AssertionError.class,
                 "Staged session " + sessionId + " already contains " + SHIM_APEX_PACKAGE_NAME,
                 Install.single(APEX_V2));
-
     }
 
+    @Test
+    public void testGetStagedModuleNames() throws Exception {
+        // Before staging a session
+        String[] result = getPackageManagerNative().getStagedApexModuleNames();
+        assertThat(result).hasLength(0);
+        // Stage an apex
+        int sessionId = Install.single(APEX_V2).setStaged().commit();
+        waitForSessionReady(sessionId);
+        result = getPackageManagerNative().getStagedApexModuleNames();
+        assertThat(result).hasLength(1);
+        assertThat(result).isEqualTo(new String[]{SHIM_APEX_PACKAGE_NAME});
+        // Abandon the session
+        InstallUtils.openPackageInstallerSession(sessionId).abandon();
+        result = getPackageManagerNative().getStagedApexModuleNames();
+        assertThat(result).hasLength(0);
+    }
+
+    @Test
+    public void testGetStagedApexInfo() throws Exception {
+        // Ask for non-existing module
+        StagedApexInfo result = getPackageManagerNative().getStagedApexInfo("not found");
+        assertThat(result).isNull();
+        // Stage an apex
+        int sessionId = Install.single(APEX_V2).setStaged().commit();
+        waitForSessionReady(sessionId);
+        // Query proper module name
+        result = getPackageManagerNative().getStagedApexInfo(SHIM_APEX_PACKAGE_NAME);
+        assertThat(result.moduleName).isEqualTo(SHIM_APEX_PACKAGE_NAME);
+        InstallUtils.openPackageInstallerSession(sessionId).abandon();
+    }
+
+    public static class MockStagedApexObserver extends IStagedApexObserver.Stub {
+        @Override
+        public void onApexStaged(ApexStagedEvent event) {
+            assertThat(event).isNotNull();
+        }
+    }
+
+    @Test
+    public void testStagedApexObserver() throws Exception {
+        MockStagedApexObserver realObserver = new MockStagedApexObserver();
+        IStagedApexObserver observer = spy(realObserver);
+        assertThat(observer).isNotNull();
+        getPackageManagerNative().registerStagedApexObserver(observer);
+
+        // Stage an apex and verify observer was called
+        int sessionId = Install.single(APEX_V2).setStaged().commit();
+        waitForSessionReady(sessionId);
+        ArgumentCaptor<ApexStagedEvent> captor = ArgumentCaptor.forClass(ApexStagedEvent.class);
+        verify(observer, timeout(5000)).onApexStaged(captor.capture());
+        assertThat(captor.getValue().stagedApexModuleNames).isEqualTo(
+                new String[] {SHIM_APEX_PACKAGE_NAME});
+
+        // Abandon and verify observer is called
+        Mockito.clearInvocations(observer);
+        InstallUtils.openPackageInstallerSession(sessionId).abandon();
+        verify(observer, timeout(5000)).onApexStaged(captor.capture());
+        assertThat(captor.getValue().stagedApexModuleNames).hasLength(0);
+    }
+
+    private IPackageManagerNative getPackageManagerNative() {
+        IBinder binder = ServiceManager.waitForService("package_native");
+        assertThat(binder).isNotNull();
+        return IPackageManagerNative.Stub.asInterface(binder);
+    }
     private static void assertSessionApplied(int sessionId) {
         assertSessionState(sessionId, (session) -> {
             assertThat(session.isStagedSessionApplied()).isTrue();

tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java

@@ -478,6 +478,21 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         runPhase("testRebootlessUpdate_hasStagedSessionWithSameApex_fails");
     }
 
+    @Test
+    public void testGetStagedModuleNames() throws Exception {
+        runPhase("testGetStagedModuleNames");
+    }
+
+    @Test
+    public void testGetStagedApexInfo() throws Exception {
+        runPhase("testGetStagedApexInfo");
+    }
+
+    @Test
+    public void testStagedApexObserver() throws Exception {
+        runPhase("testStagedApexObserver");
+    }
+
     private List<String> getStagingDirectories() throws DeviceNotAvailableException {
         String baseDir = "/data/app-staging";
         try {