Add unit tests for SplitAttributes runtime APIs

Test: atest WMJetpackUnitTests
Bug: 236764532

Change-Id: If46bf1ca9e1db7848655dd3d26a88b64660f2874
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 2c1ddf7..9f3c19b 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -308,7 +308,7 @@
                 synchronized (mLock) {
                     final List<TaskFragmentContainer> containers = taskContainer.mContainers;
                     // Clean up the TaskFragmentContainers by the z-order from the lowest.
-                    for (int i = 0; i < containers.size() - 1; i++) {
+                    for (int i = 0; i < containers.size(); i++) {
                         final TaskFragmentContainer container = containers.get(i);
                         if (pendingFinishingContainers.contains(container)) {
                             // Don't update records here to prevent double invocation.
@@ -318,7 +318,7 @@
                     }
                     // Remove container records.
                     removeContainers(taskContainer, pendingFinishingContainers);
-                    // Update the change to the client side.
+                    // Update the change to the server side.
                     updateContainersInTaskIfVisible(wct, taskContainer.getTaskId());
                 }
             });
@@ -353,21 +353,25 @@
     @Override
     public void updateSplitAttributes(@NonNull IBinder splitInfoToken,
             @NonNull SplitAttributes splitAttributes) {
+        Objects.requireNonNull(splitInfoToken);
+        Objects.requireNonNull(splitAttributes);
         synchronized (mLock) {
             final SplitContainer splitContainer = getSplitContainer(splitInfoToken);
             if (splitContainer == null) {
                 Log.w(TAG, "Cannot find SplitContainer for token:" + splitInfoToken);
                 return;
             }
-            WindowContainerTransaction wct = mTransactionManager.startNewTransaction()
-                    .getTransaction();
+            // Override the default split Attributes so that it will be applied
+            // if the SplitContainer is not visible currently.
+            splitContainer.updateDefaultSplitAttributes(splitAttributes);
+
+            final TransactionRecord transactionRecord = mTransactionManager.startNewTransaction();
+            final WindowContainerTransaction wct = transactionRecord.getTransaction();
             if (updateSplitContainerIfNeeded(splitContainer, wct, splitAttributes)) {
-                splitContainer.updateDefaultSplitAttributes(splitAttributes);
-                mTransactionManager.getCurrentTransactionRecord()
-                        .apply(false /* shouldApplyIndependently */);
+                transactionRecord.apply(false /* shouldApplyIndependently */);
             } else {
                 // Abort if the SplitContainer wasn't updated.
-                mTransactionManager.getCurrentTransactionRecord().abort();
+                transactionRecord.abort();
             }
         }
     }
@@ -1559,8 +1563,9 @@
      *
      * @return {@code true} if the update succeed. Otherwise, returns {@code false}.
      */
+    @VisibleForTesting
     @GuardedBy("mLock")
-    private boolean updateSplitContainerIfNeeded(@NonNull SplitContainer splitContainer,
+    boolean updateSplitContainerIfNeeded(@NonNull SplitContainer splitContainer,
             @NonNull WindowContainerTransaction wct, @Nullable SplitAttributes splitAttributes) {
         if (!isTopMostSplit(splitContainer)) {
             // Skip position update - it isn't the topmost split.
@@ -1904,6 +1909,7 @@
         return null;
     }
 
+    @VisibleForTesting
     @Nullable
     @GuardedBy("mLock")
     SplitContainer getSplitContainer(@NonNull IBinder token) {
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
index 17909d4..61cd748 100644
--- 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
@@ -67,6 +67,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 
 import android.annotation.NonNull;
 import android.app.Activity;
@@ -82,6 +83,7 @@
 import android.os.Handler;
 import android.os.IBinder;
 import android.platform.test.annotations.Presubmit;
+import android.util.ArraySet;
 import android.view.WindowInsets;
 import android.view.WindowMetrics;
 import android.window.TaskFragmentInfo;
@@ -100,12 +102,14 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Consumer;
 
 /**
@@ -1323,6 +1327,125 @@
         verify(mTransaction).startActivityInTaskFragment(any(), any(), any(), any());
     }
 
+    @Test
+    public void testFinishActivityStacks_emptySet_earlyReturn() {
+        mSplitController.finishActivityStacks(Collections.emptySet());
+
+        verify(mSplitController, never()).updateContainersInTaskIfVisible(any(), anyInt());
+    }
+
+    @Test
+    public void testFinishActivityStacks_invalidStacks_earlyReturn() {
+        mSplitController.finishActivityStacks(Collections.singleton(new Binder()));
+
+        verify(mSplitController, never()).updateContainersInTaskIfVisible(any(), anyInt());
+    }
+
+    @Test
+    public void testFinishActivityStacks_finishSingleActivityStack() {
+        TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        tf.setInfo(mTransaction, createMockTaskFragmentInfo(tf, mActivity));
+
+        List<TaskFragmentContainer> containers = mSplitController.mTaskContainers.get(TASK_ID)
+                .mContainers;
+
+        assertEquals(containers.get(0), tf);
+
+        mSplitController.finishActivityStacks(Collections.singleton(tf.getTaskFragmentToken()));
+
+        verify(mSplitPresenter).deleteTaskFragment(any(), eq(tf.getTaskFragmentToken()));
+        assertTrue(containers.isEmpty());
+    }
+
+    @Test
+    public void testFinishActivityStacks_finishActivityStacksInOrder() {
+        TaskFragmentContainer bottomTf = mSplitController.newContainer(mActivity, TASK_ID);
+        TaskFragmentContainer topTf = mSplitController.newContainer(mActivity, TASK_ID);
+        bottomTf.setInfo(mTransaction, createMockTaskFragmentInfo(bottomTf, mActivity));
+        topTf.setInfo(mTransaction, createMockTaskFragmentInfo(topTf, createMockActivity()));
+
+        List<TaskFragmentContainer> containers = mSplitController.mTaskContainers.get(TASK_ID)
+                .mContainers;
+
+        assertEquals(containers.size(), 2);
+
+        Set<IBinder> activityStackTokens = new ArraySet<>(new IBinder[]{
+                topTf.getTaskFragmentToken(), bottomTf.getTaskFragmentToken()});
+
+        mSplitController.finishActivityStacks(activityStackTokens);
+
+        ArgumentCaptor<IBinder> argumentCaptor = ArgumentCaptor.forClass(IBinder.class);
+
+        verify(mSplitPresenter, times(2)).deleteTaskFragment(any(), argumentCaptor.capture());
+
+        List<IBinder> fragmentTokens = argumentCaptor.getAllValues();
+        assertEquals("The ActivityStack must be deleted from the lowest z-order "
+                + "regardless of the order in ActivityStack set",
+                bottomTf.getTaskFragmentToken(), fragmentTokens.get(0));
+        assertEquals("The ActivityStack must be deleted from the lowest z-order "
+                        + "regardless of the order in ActivityStack set",
+                topTf.getTaskFragmentToken(), fragmentTokens.get(1));
+
+        assertTrue(containers.isEmpty());
+    }
+
+    @Test
+    public void testUpdateSplitAttributes_invalidSplitContainerToken_earlyReturn() {
+        mSplitController.updateSplitAttributes(new Binder(), SPLIT_ATTRIBUTES);
+
+        verify(mTransactionManager, never()).startNewTransaction();
+    }
+
+    @Test
+    public void testUpdateSplitAttributes_nullParams_throwException() {
+        assertThrows(NullPointerException.class,
+                () -> mSplitController.updateSplitAttributes(null, SPLIT_ATTRIBUTES));
+
+        final SplitContainer splitContainer = mock(SplitContainer.class);
+        final IBinder token = new Binder();
+        doReturn(token).when(splitContainer).getToken();
+        doReturn(splitContainer).when(mSplitController).getSplitContainer(eq(token));
+
+        assertThrows(NullPointerException.class,
+                () -> mSplitController.updateSplitAttributes(token, null));
+    }
+
+    @Test
+    public void testUpdateSplitAttributes_doNotNeedToUpdateSplitContainer_doNotApplyTransaction() {
+        final SplitContainer splitContainer = mock(SplitContainer.class);
+        final IBinder token = new Binder();
+        doReturn(token).when(splitContainer).getToken();
+        doReturn(splitContainer).when(mSplitController).getSplitContainer(eq(token));
+        doReturn(false).when(mSplitController).updateSplitContainerIfNeeded(
+                eq(splitContainer), any(), eq(SPLIT_ATTRIBUTES));
+        TransactionManager.TransactionRecord testRecord =
+                mock(TransactionManager.TransactionRecord.class);
+        doReturn(testRecord).when(mTransactionManager).startNewTransaction();
+
+        mSplitController.updateSplitAttributes(token, SPLIT_ATTRIBUTES);
+
+        verify(splitContainer).updateDefaultSplitAttributes(eq(SPLIT_ATTRIBUTES));
+        verify(testRecord).abort();
+    }
+
+    @Test
+    public void testUpdateSplitAttributes_splitContainerUpdated_updateAttrs() {
+        final SplitContainer splitContainer = mock(SplitContainer.class);
+        final IBinder token = new Binder();
+        doReturn(token).when(splitContainer).getToken();
+        doReturn(splitContainer).when(mSplitController).getSplitContainer(eq(token));
+        doReturn(true).when(mSplitController).updateSplitContainerIfNeeded(
+                eq(splitContainer), any(), eq(SPLIT_ATTRIBUTES));
+        TransactionManager.TransactionRecord testRecord =
+                mock(TransactionManager.TransactionRecord.class);
+        doReturn(testRecord).when(mTransactionManager).startNewTransaction();
+
+        mSplitController.updateSplitAttributes(token, SPLIT_ATTRIBUTES);
+
+        verify(splitContainer).updateDefaultSplitAttributes(eq(SPLIT_ATTRIBUTES));
+        verify(testRecord).apply(eq(false));
+    }
+
     /** Creates a mock activity in the organizer process. */
     private Activity createMockActivity() {
         return createMockActivity(TASK_ID);