Remove no cards card after undo before adding the card back in.

PiperOrigin-RevId: 252112494
Change-Id: I3c094482314384dbfe7260f859a5075d8e70e480
diff --git a/src/main/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriver.java b/src/main/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriver.java
index 88f8fdb..284051f 100644
--- a/src/main/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriver.java
+++ b/src/main/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriver.java
@@ -388,6 +388,16 @@
     maybeRestoreScroll();
   }
 
+  private void alwaysRemoveNoContentOrZeroStateCardIfPresent() {
+    if (featureDrivers.size() == 1
+        && ((featureDrivers.get(0) instanceof NoContentDriver)
+            || (featureDrivers.get(0) instanceof ZeroStateDriver))) {
+      featureDrivers.get(0).onDestroy();
+      featureDrivers.remove(0);
+      notifyContentRemoved(0);
+    }
+  }
+
   private void maybeRemoveNoContentOrZeroStateCard() {
     if (shouldRemoveNoContentCardOrZeroState()) {
       featureDrivers.get(0).onDestroy();
@@ -677,7 +687,7 @@
 
           @Override
           public void onDismissedWithAction() {
-            maybeRemoveNoContentOrZeroStateCard();
+            alwaysRemoveNoContentOrZeroStateCardIfPresent();
             pendingDismissCallback.onDismissReverted();
             createFeatureChildAndInsertAtIndex(modelChild, index);
           }
diff --git a/src/test/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriverTest.java b/src/test/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriverTest.java
index 91e4ff8..f575efe 100644
--- a/src/test/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriverTest.java
+++ b/src/test/java/com/google/android/libraries/feed/basicstream/internal/drivers/StreamDriverTest.java
@@ -17,7 +17,6 @@
 import static com.google.android.libraries.feed.common.testing.RunnableSubject.assertThatRunnable;
 import static com.google.android.libraries.feed.testing.modelprovider.FakeModelCursor.getCardModelFeatureWithCursor;
 import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -91,7 +90,7 @@
 /** Tests for {@link StreamDriver}. */
 @RunWith(RobolectricTestRunner.class)
 public class StreamDriverTest {
-
+  private static final String CONTENT_ID = "contentID";
   private static final FakeModelCursor CURSOR_WITH_INVALID_CHILDREN =
       FakeModelCursor.newBuilder()
           .addChild(
@@ -782,35 +781,63 @@
 
   @Test
   public void testTriggerPendingDismiss_noAction() {
+    setCards(CONTENT_ID);
     SnackbarCallbackApi snackbarCallbackApi =
-        triggerPendingDismiss(UndoAction.newBuilder().setConfirmationLabel("confirmation").build());
+        triggerPendingDismiss(
+            UndoAction.newBuilder().setConfirmationLabel("confirmation").build(), CONTENT_ID);
 
     assertThat(streamDriver.isZeroStateBeingShown()).isTrue();
 
     snackbarCallbackApi.onDismissNoAction();
-    
+
     verify(pendingDismissCallback).onDismissCommitted();
     assertThat(streamDriver.isZeroStateBeingShown()).isTrue();
   }
 
   @Test
-  public void testTriggerPendingDismiss_actionTaken() {
+  public void testTriggerPendingDismiss_actionTaken_lastCard() {
+    setCards(CONTENT_ID);
     SnackbarCallbackApi snackbarCallbackApi =
-        triggerPendingDismiss(UndoAction.newBuilder().setConfirmationLabel("confirmation").build());
+        triggerPendingDismiss(
+            UndoAction.newBuilder().setConfirmationLabel("confirmation").build(), CONTENT_ID);
 
+    List<LeafFeatureDriver> leafFeatureDrivers = streamDriver.getLeafFeatureDrivers();
+    assertThat(leafFeatureDrivers).hasSize(1);
     assertThat(streamDriver.isZeroStateBeingShown()).isTrue();
 
     snackbarCallbackApi.onDismissedWithAction();
+
+    verify(pendingDismissCallback).onDismissReverted();
+    leafFeatureDrivers = streamDriver.getLeafFeatureDrivers();
+    assertThat(leafFeatureDrivers).hasSize(1);
+    assertThat(streamDriver.isZeroStateBeingShown()).isFalse();
+  }
+
+  @Test
+  public void testTriggerPendingDismiss_actionTaken_multipleCards() {
+    setCards(CONTENT_ID, CONTENT_ID + "1");
+    SnackbarCallbackApi snackbarCallbackApi =
+        triggerPendingDismiss(
+            UndoAction.newBuilder().setConfirmationLabel("confirmation").build(), CONTENT_ID);
+
+    List<LeafFeatureDriver> leafFeatureDrivers = streamDriver.getLeafFeatureDrivers();
+    assertThat(leafFeatureDrivers).hasSize(1);
+    assertThat(streamDriver.isZeroStateBeingShown()).isFalse();
+
+    snackbarCallbackApi.onDismissedWithAction();
     
     verify(pendingDismissCallback).onDismissReverted();
-    verify(contentListener, times(2)).notifyContentsAdded(eq(0), any());
+    leafFeatureDrivers = streamDriver.getLeafFeatureDrivers();
+    assertThat(leafFeatureDrivers).hasSize(2);
     assertThat(streamDriver.isZeroStateBeingShown()).isFalse();
   }
 
   @Test
   public void testTriggerPendingDismiss_actionStringSent() {
+    setCards(CONTENT_ID);
     triggerPendingDismiss(
-        UndoAction.newBuilder().setConfirmationLabel("conf").setUndoLabel("undo label").build());
+        UndoAction.newBuilder().setConfirmationLabel("conf").setUndoLabel("undo label").build(),
+        CONTENT_ID);
   }
 
   @Test
@@ -934,19 +961,22 @@
     return streamDriver.childrenCreated.get(i).getLeafFeatureDriver();
   }
 
-  private SnackbarCallbackApi triggerPendingDismiss(UndoAction undoAction) {
+  private void setCards(String... contentId) {
+    FakeModelCursor.Builder cursor = FakeModelCursor.newBuilder();
+    for (String id : contentId) {
+      cursor.addChild(
+          FakeModelChild.newBuilder()
+              .setModelFeature(getCardModelFeatureWithCursor(EMPTY_MODEL_CURSOR))
+              .setContentId(id)
+              .build());
+    }
+    when(streamFeature.getCursor()).thenReturn(cursor.build());
+  }
+
+  private SnackbarCallbackApi triggerPendingDismiss(UndoAction undoAction, String contentId) {
     ArgumentCaptor<SnackbarCallbackApi> snackbarCallbackApi =
         ArgumentCaptor.forClass(SnackbarCallbackApi.class);
-    String contentId = "contentId";
-    when(streamFeature.getCursor())
-        .thenReturn(
-            FakeModelCursor.newBuilder()
-                .addChild(
-                    FakeModelChild.newBuilder()
-                        .setModelFeature(getCardModelFeatureWithCursor(EMPTY_MODEL_CURSOR))
-                        .setContentId(contentId)
-                        .build())
-                .build());
+
 
     // Causes StreamDriver to build a list of children based on the children from the cursor.
     streamDriver.getLeafFeatureDrivers();
@@ -1055,5 +1085,4 @@
       return zeroStateDriver;
     }
   }
-
 }