diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 418ff0e..bb3b534 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -43,6 +43,8 @@
 
 import androidx.window.common.EmptyLifecycleCallbacksAdapter;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
@@ -65,10 +67,12 @@
      * When the app is host of multiple Tasks, there can be multiple splits controlled by the same
      * organizer.
      */
-    private final SparseArray<TaskContainer> mTaskContainers = new SparseArray<>();
+    @VisibleForTesting
+    final SparseArray<TaskContainer> mTaskContainers = new SparseArray<>();
 
     // Callback to Jetpack to notify about changes to split states.
-    private @NonNull Consumer<List<SplitInfo>> mEmbeddingCallback;
+    @NonNull
+    private Consumer<List<SplitInfo>> mEmbeddingCallback;
     private final List<SplitInfo> mLastReportedSplitStates = new ArrayList<>();
 
     // We currently only support split activity embedding within the one root Task.
@@ -1029,7 +1033,8 @@
     }
 
     /** Represents TaskFragments and split pairs below a Task. */
-    private static class TaskContainer {
+    @VisibleForTesting
+    static class TaskContainer {
         final List<TaskFragmentContainer> mContainers = new ArrayList<>();
         final List<SplitContainer> mSplitContainers = new ArrayList<>();
     }
diff --git a/libs/WindowManager/Jetpack/tests/unittest/Android.bp b/libs/WindowManager/Jetpack/tests/unittest/Android.bp
index 62e8128..212fbd0 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/Android.bp
+++ b/libs/WindowManager/Jetpack/tests/unittest/Android.bp
@@ -46,6 +46,12 @@
         "android.test.runner",
     ],
 
+    // These are not normally accessible from apps so they must be explicitly included.
+    jni_libs: [
+        "libdexmakerjvmtiagent",
+        "libstaticjvmtiagent",
+    ],
+
     optimize: {
         enabled: false,
     },
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
new file mode 100644
index 0000000..30e89a6
--- /dev/null
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.window.extensions.embedding;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+import androidx.window.extensions.embedding.SplitController.TaskContainer;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class SplitControllerTest {
+    private static final int TASK_ID = 10;
+
+    private SplitController mSplitController;
+
+    @Before
+    public void setUp() {
+        mSplitController = new SplitController();
+        spyOn(mSplitController);
+    }
+
+    @Test
+    public void testGetTopActiveContainer() {
+        TaskContainer taskContainer = new TaskContainer();
+        // tf3 is finished so is not active.
+        TaskFragmentContainer tf3 = mock(TaskFragmentContainer.class);
+        doReturn(true).when(tf3).isFinished();
+        // tf2 has running activity so is active.
+        TaskFragmentContainer tf2 = mock(TaskFragmentContainer.class);
+        doReturn(1).when(tf2).getRunningActivityCount();
+        // tf1 has no running activity so is not active.
+        TaskFragmentContainer tf1 = new TaskFragmentContainer(null, TASK_ID);
+
+        taskContainer.mContainers.add(tf3);
+        taskContainer.mContainers.add(tf2);
+        taskContainer.mContainers.add(tf1);
+        mSplitController.mTaskContainers.put(TASK_ID, taskContainer);
+
+        assertWithMessage("Must return tf2 because tf3 is not active.")
+                .that(mSplitController.getTopActiveContainer(TASK_ID)).isEqualTo(tf2);
+
+        taskContainer.mContainers.remove(tf1);
+
+        assertWithMessage("Must return tf2 because tf2 has running activity.")
+                .that(mSplitController.getTopActiveContainer(TASK_ID)).isEqualTo(tf2);
+
+        taskContainer.mContainers.remove(tf2);
+
+        assertWithMessage("Must return null because tf1 has no running activity.")
+                .that(mSplitController.getTopActiveContainer(TASK_ID)).isNull();
+    }
+}