diff --git a/core/java/android/app/prediction/AppPredictionContext.java b/core/java/android/app/prediction/AppPredictionContext.java
index 87ccb66..2fbe6e36 100644
--- a/core/java/android/app/prediction/AppPredictionContext.java
+++ b/core/java/android/app/prediction/AppPredictionContext.java
@@ -18,6 +18,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.content.Context;
 import android.os.Bundle;
 import android.os.Parcel;
@@ -29,6 +30,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppPredictionContext implements Parcelable {
 
     private final int mPredictedTargetCount;
@@ -73,6 +75,17 @@
     }
 
     @Override
+    public boolean equals(Object o) {
+        if (o == this) return true;
+        if (!getClass().equals(o != null ? o.getClass() : null)) return false;
+
+        AppPredictionContext other = (AppPredictionContext) o;
+        return mPredictedTargetCount == other.mPredictedTargetCount
+                && mUiSurface.equals(other.mUiSurface)
+                && mPackageName.equals(other.mPackageName);
+    }
+
+    @Override
     public int describeContents() {
         return 0;
     }
@@ -104,6 +117,7 @@
      * @hide
      */
     @SystemApi
+    @TestApi
     public static final class Builder {
 
         @NonNull
@@ -116,8 +130,12 @@
         private Bundle mExtras;
 
         /**
+         * TODO(b/123591863): Add java docs
+         *
          * @hide
          */
+        @SystemApi
+        @TestApi
         public Builder(@NonNull Context context) {
             mPackageName = context.getPackageName();
         }
diff --git a/core/java/android/app/prediction/AppPredictionManager.java b/core/java/android/app/prediction/AppPredictionManager.java
index f8578d4..99f78f1 100644
--- a/core/java/android/app/prediction/AppPredictionManager.java
+++ b/core/java/android/app/prediction/AppPredictionManager.java
@@ -17,6 +17,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.content.Context;
 
 import com.android.internal.util.Preconditions;
@@ -26,6 +27,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppPredictionManager {
 
     private final Context mContext;
diff --git a/core/java/android/app/prediction/AppPredictionSessionId.java b/core/java/android/app/prediction/AppPredictionSessionId.java
index 1d7308e..1e76c24 100644
--- a/core/java/android/app/prediction/AppPredictionSessionId.java
+++ b/core/java/android/app/prediction/AppPredictionSessionId.java
@@ -17,6 +17,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -26,6 +27,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppPredictionSessionId implements Parcelable {
 
     private final String mId;
diff --git a/core/java/android/app/prediction/AppPredictor.java b/core/java/android/app/prediction/AppPredictor.java
index 2ddbd08..12d6ce3 100644
--- a/core/java/android/app/prediction/AppPredictor.java
+++ b/core/java/android/app/prediction/AppPredictor.java
@@ -19,6 +19,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.app.prediction.IPredictionCallback.Stub;
 import android.content.Context;
 import android.content.pm.ParceledListSlice;
@@ -63,6 +64,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppPredictor {
 
     private static final String TAG = AppPredictor.class.getSimpleName();
@@ -102,6 +104,10 @@
      * Notifies the prediction service of an app target event.
      */
     public void notifyAppTargetEvent(@NonNull AppTargetEvent event) {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         try {
             mPredictionManager.notifyAppTargetEvent(mSessionId, event);
         } catch (RemoteException e) {
@@ -114,6 +120,10 @@
      */
     public void notifyLocationShown(@NonNull String launchLocation,
             @NonNull List<AppTargetId> targetIds) {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         try {
             mPredictionManager.notifyLocationShown(mSessionId, launchLocation,
                     new ParceledListSlice<>(targetIds));
@@ -130,6 +140,10 @@
      */
     public void registerPredictionUpdates(@NonNull @CallbackExecutor Executor callbackExecutor,
             @NonNull AppPredictor.Callback callback) {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         if (mRegisteredCallbacks.containsKey(callback)) {
             // Skip if this callback is already registered
             return;
@@ -149,6 +163,10 @@
      * callback until the callback is re-registered.
      */
     public void unregisterPredictionUpdates(@NonNull AppPredictor.Callback callback) {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         if (!mRegisteredCallbacks.containsKey(callback)) {
             // Skip if this callback was never registered
             return;
@@ -168,6 +186,10 @@
      * @see Callback#onTargetsAvailable(List)
      */
     public void requestPredictionUpdate() {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         try {
             mPredictionManager.requestPredictionUpdate(mSessionId);
         } catch (RemoteException e) {
@@ -182,6 +204,10 @@
     @Nullable
     public void sortTargets(@NonNull List<AppTarget> targets,
             @NonNull Executor callbackExecutor, @NonNull Consumer<List<AppTarget>> callback) {
+        if (mIsClosed.get()) {
+            throw new IllegalStateException("This client has already been destroyed.");
+        }
+
         try {
             mPredictionManager.sortAppTargets(mSessionId, new ParceledListSlice(targets),
                     new CallbackWrapper(callbackExecutor, callback));
@@ -206,6 +232,8 @@
             } catch (RemoteException e) {
                 Log.e(TAG, "Failed to notify app target event", e);
             }
+        } else {
+            throw new IllegalStateException("This client has already been destroyed.");
         }
     }
 
@@ -222,6 +250,16 @@
     }
 
     /**
+     * TODO(b/123591863): Add java docs
+     *
+     * @hide
+     */
+    @TestApi
+    public AppPredictionSessionId getSessionId() {
+        return mSessionId;
+    }
+
+    /**
      * Callback for receiving prediction updates.
      */
     public interface Callback {
diff --git a/core/java/android/app/prediction/AppTarget.java b/core/java/android/app/prediction/AppTarget.java
index 99c1c44..b924cec 100644
--- a/core/java/android/app/prediction/AppTarget.java
+++ b/core/java/android/app/prediction/AppTarget.java
@@ -18,6 +18,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.content.pm.ShortcutInfo;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -30,6 +31,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppTarget implements Parcelable {
 
     private final AppTargetId mId;
@@ -42,8 +44,12 @@
     private int mRank;
 
     /**
+     * TODO(b/123591863): Add java docs
+     *
      * @hide
      */
+    @SystemApi
+    @TestApi
     public AppTarget(@NonNull AppTargetId id, @NonNull String packageName,
             @Nullable String className, @NonNull UserHandle user) {
         mId = id;
@@ -55,15 +61,19 @@
     }
 
     /**
+     * TODO(b/123591863): Add java docs
+     *
      * @hide
      */
-    public AppTarget(@NonNull AppTargetId id, @NonNull ShortcutInfo shortcutInfo) {
+    @SystemApi
+    public AppTarget(@NonNull AppTargetId id, @NonNull ShortcutInfo shortcutInfo,
+            @Nullable String className) {
         mId = id;
         mShortcutInfo = Preconditions.checkNotNull(shortcutInfo);
 
         mPackageName = mShortcutInfo.getPackage();
         mUser = mShortcutInfo.getUserHandle();
-        mClassName = null;
+        mClassName = className;
     }
 
     private AppTarget(Parcel parcel) {
@@ -71,13 +81,12 @@
         mShortcutInfo = parcel.readTypedObject(ShortcutInfo.CREATOR);
         if (mShortcutInfo == null) {
             mPackageName = parcel.readString();
-            mClassName = parcel.readString();
             mUser = UserHandle.of(parcel.readInt());
         } else {
             mPackageName = mShortcutInfo.getPackage();
             mUser = mShortcutInfo.getUserHandle();
-            mClassName = null;
         }
+        mClassName = parcel.readString();
         mRank = parcel.readInt();
     }
 
@@ -129,11 +138,32 @@
         mRank = rank;
     }
 
+    /**
+     * Returns the rank for the target.
+     */
     public int getRank() {
         return mRank;
     }
 
     @Override
+    public boolean equals(Object o) {
+        if (!getClass().equals(o != null ? o.getClass() : null)) return false;
+
+        AppTarget other = (AppTarget) o;
+        boolean sameClassName = (mClassName == null && other.mClassName == null)
+                || (mClassName != null && mClassName.equals(other.mClassName));
+        boolean sameShortcutInfo = (mShortcutInfo == null && other.mShortcutInfo == null)
+                || (mShortcutInfo != null && other.mShortcutInfo != null
+                        && mShortcutInfo.getId() == other.mShortcutInfo.getId());
+        return mId.equals(other.mId)
+                && mPackageName.equals(other.mPackageName)
+                && sameClassName
+                && mUser.equals(other.mUser)
+                && sameShortcutInfo
+                && mRank == other.mRank;
+    }
+
+    @Override
     public int describeContents() {
         return 0;
     }
@@ -144,9 +174,9 @@
         dest.writeTypedObject(mShortcutInfo, flags);
         if (mShortcutInfo == null) {
             dest.writeString(mPackageName);
-            dest.writeString(mClassName);
             dest.writeInt(mUser.getIdentifier());
         }
+        dest.writeString(mClassName);
         dest.writeInt(mRank);
     }
 
diff --git a/core/java/android/app/prediction/AppTargetEvent.java b/core/java/android/app/prediction/AppTargetEvent.java
index 18317e1..37e41de 100644
--- a/core/java/android/app/prediction/AppTargetEvent.java
+++ b/core/java/android/app/prediction/AppTargetEvent.java
@@ -19,6 +19,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -30,6 +31,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppTargetEvent implements Parcelable {
 
     /**
@@ -96,6 +98,16 @@
     }
 
     @Override
+    public boolean equals(Object o) {
+        if (!getClass().equals(o != null ? o.getClass() : null)) return false;
+
+        AppTargetEvent other = (AppTargetEvent) o;
+        return mTarget.equals(other.mTarget)
+                && mLocation.equals(other.mLocation)
+                && mAction == other.mAction;
+    }
+
+    @Override
     public int describeContents() {
         return 0;
     }
@@ -126,6 +138,7 @@
      * @hide
      */
     @SystemApi
+    @TestApi
     public static final class Builder {
         private AppTarget mTarget;
         private String mLocation;
diff --git a/core/java/android/app/prediction/AppTargetId.java b/core/java/android/app/prediction/AppTargetId.java
index 0b8fb47..639ba78 100644
--- a/core/java/android/app/prediction/AppTargetId.java
+++ b/core/java/android/app/prediction/AppTargetId.java
@@ -17,6 +17,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -25,14 +26,19 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public final class AppTargetId implements Parcelable {
 
     @NonNull
     private final String mId;
 
     /**
+     * TODO(b/123591863): Add java docs
+     *
      * @hide
      */
+    @SystemApi
+    @TestApi
     public AppTargetId(@NonNull String id) {
         mId = id;
     }
diff --git a/core/java/android/service/appprediction/AppPredictionService.java b/core/java/android/service/appprediction/AppPredictionService.java
index b77405a..d012851 100644
--- a/core/java/android/service/appprediction/AppPredictionService.java
+++ b/core/java/android/service/appprediction/AppPredictionService.java
@@ -21,6 +21,7 @@
 import android.annotation.MainThread;
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.app.Service;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppPredictionSessionId;
@@ -49,6 +50,7 @@
  * @hide
  */
 @SystemApi
+@TestApi
 public abstract class AppPredictionService extends Service {
 
     private static final String TAG = "AppPredictionService";
@@ -140,6 +142,7 @@
 
     @Override
     public final IBinder onBind(Intent intent) {
+        // TODO(b/111701043): Verify that the action is valid
         return mInterface.asBinder();
     }
 
@@ -228,6 +231,7 @@
     public void onStopPredictionUpdates() {}
 
     private void doRequestPredictionUpdate(@NonNull AppPredictionSessionId sessionId) {
+        // Just an optimization, if there are no callbacks, then don't bother notifying the service
         final ArrayList<CallbackWrapper> callbacks = mSessionCallbacks.get(sessionId);
         if (callbacks != null && !callbacks.isEmpty()) {
             onRequestPredictionUpdate(sessionId);