Refactor GcTest to include PietSharedState objects.

PiperOrigin-RevId: 251446714
Change-Id: Ib457acb92bf3589ab1f91ab60ac4b22448657d88
diff --git a/src/main/java/com/google/android/libraries/feed/common/testing/BUILD b/src/main/java/com/google/android/libraries/feed/common/testing/BUILD
index 963e4a2..691b992 100644
--- a/src/main/java/com/google/android/libraries/feed/common/testing/BUILD
+++ b/src/main/java/com/google/android/libraries/feed/common/testing/BUILD
@@ -39,6 +39,7 @@
         "//src/main/java/com/google/android/libraries/feed/testing/requestmanager",
         "//src/main/proto/com/google/android/libraries/feed/api/internal/proto:client_feed_java_proto_lite",
         "//src/main/proto/search/now/ui/piet:piet_java_proto_lite",
+        "//src/main/proto/search/now/ui/stream:stream_java_proto_lite",
         "//src/main/proto/search/now/wire/feed:feed_java_proto_lite",
         "@com_google_code_findbugs_jsr305//jar",
         "@com_google_protobuf_javalite//:protobuf_java_lite",
diff --git a/src/main/java/com/google/android/libraries/feed/common/testing/ResponseBuilder.java b/src/main/java/com/google/android/libraries/feed/common/testing/ResponseBuilder.java
index 51c192c..fb28e5d 100644
--- a/src/main/java/com/google/android/libraries/feed/common/testing/ResponseBuilder.java
+++ b/src/main/java/com/google/android/libraries/feed/common/testing/ResponseBuilder.java
@@ -16,6 +16,8 @@
 
 import com.google.protobuf.ByteString;
 import com.google.search.now.ui.piet.PietProto.PietSharedState;
+import com.google.search.now.ui.stream.StreamStructureProto.Content;
+import com.google.search.now.ui.stream.StreamStructureProto.PietContent;
 import com.google.search.now.wire.feed.ActionPropertiesProto.ActionProperties;
 import com.google.search.now.wire.feed.ContentIdProto.ContentId;
 import com.google.search.now.wire.feed.DataOperationProto.DataOperation;
@@ -52,6 +54,7 @@
 
   private final WireProtocolInfo wireProtocolInfo = new WireProtocolInfo();
   private final FeedResponse.Builder feedResponseBuilder = FeedResponse.newBuilder();
+  private final List<ContentId> pietSharedStateContentIds = new ArrayList<>();
 
   private int tokenId = 0;
   /*@Nullable*/ private ByteString token = null;
@@ -72,8 +75,13 @@
   }
 
   public ResponseBuilder addPietSharedState() {
+    return addPietSharedState(PIET_SHARED_STATE);
+  }
+
+  public ResponseBuilder addPietSharedState(ContentId pietSharedStateContentId) {
+    pietSharedStateContentIds.add(pietSharedStateContentId);
     PayloadMetadata payloadMetadata =
-        PayloadMetadata.newBuilder().setContentId(PIET_SHARED_STATE).build();
+        PayloadMetadata.newBuilder().setContentId(pietSharedStateContentId).build();
     feedResponseBuilder.addDataOperation(
         DataOperation.newBuilder()
             .setOperation(Operation.UPDATE_OR_APPEND)
@@ -141,6 +149,16 @@
         Feature.newBuilder()
             .setRenderableUnit(RenderableUnit.CONTENT)
             .setParentId(contentId)
+            .setExtension(
+                Content.contentExtension,
+                Content.newBuilder()
+                    .setType(Content.Type.PIET)
+                    .setExtension(
+                        PietContent.pietContentExtension,
+                        PietContent.newBuilder()
+                            .addAllPietSharedStates(pietSharedStateContentIds)
+                            .build())
+                    .build())
             .build();
     feedResponseBuilder.addDataOperation(
         DataOperation.newBuilder()
diff --git a/src/test/java/com/google/android/libraries/feed/infraintegration/GcTest.java b/src/test/java/com/google/android/libraries/feed/infraintegration/GcTest.java
index 4e595e9..708c1ae 100644
--- a/src/test/java/com/google/android/libraries/feed/infraintegration/GcTest.java
+++ b/src/test/java/com/google/android/libraries/feed/infraintegration/GcTest.java
@@ -24,10 +24,12 @@
 import com.google.android.libraries.feed.common.testing.InfraIntegrationScope;
 import com.google.android.libraries.feed.common.testing.ResponseBuilder;
 import com.google.android.libraries.feed.common.time.testing.FakeClock;
+import com.google.search.now.feed.client.StreamDataProto.StreamSharedState;
 import com.google.search.now.feed.client.StreamDataProto.UiContext;
 import com.google.search.now.wire.feed.ContentIdProto.ContentId;
 import com.google.search.now.wire.feed.ResponseProto.Response;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.Test;
@@ -37,6 +39,18 @@
 /** Tests that assert the behavior of garbage collection. */
 @RunWith(RobolectricTestRunner.class)
 public final class GcTest {
+  private static final ContentId PIET_SHARED_STATE_1 =
+      ContentId.newBuilder()
+          .setContentDomain("piet-shared-state")
+          .setId(1)
+          .setTable("feature")
+          .build();
+  private static final ContentId PIET_SHARED_STATE_2 =
+      ContentId.newBuilder()
+          .setContentDomain("piet-shared-state")
+          .setId(2)
+          .setTable("feature")
+          .build();
   private static final ContentId[] REQUEST_1 =
       new ContentId[] {
         ResponseBuilder.createFeatureContentId(1), ResponseBuilder.createFeatureContentId(2)
@@ -58,7 +72,7 @@
   public void testGc_contentInLiveSessionRetained() {
     scope
         .getFakeFeedRequestManager()
-        .queueResponse(createResponse(REQUEST_1))
+        .queueResponse(createResponse(REQUEST_1, PIET_SHARED_STATE_1))
         .triggerRefresh(
             RequestReason.OPEN_WITHOUT_CONTENT,
             scope.getFeedSessionManager().getUpdateConsumer(MutationContext.EMPTY_CONTEXT));
@@ -73,7 +87,7 @@
     // Populate HEAD with new data.
     scope
         .getFakeFeedRequestManager()
-        .queueResponse(createResponse(REQUEST_2))
+        .queueResponse(createResponse(REQUEST_2, PIET_SHARED_STATE_2))
         .triggerRefresh(
             RequestReason.OPEN_WITHOUT_CONTENT,
             scope.getFeedSessionManager().getUpdateConsumer(MutationContext.EMPTY_CONTEXT));
@@ -82,14 +96,16 @@
     fakeClock.advance(LIFETIME_MS / 2);
     InfraIntegrationScope secondScope = scope.clone();
     assertPayloads(REQUEST_1, secondScope, /* shouldExist= */ true);
+    assertSharedStates(new ContentId[] {PIET_SHARED_STATE_1}, secondScope, /* shouldExist= */ true);
     assertPayloads(REQUEST_2, secondScope, /* shouldExist= */ true);
+    assertSharedStates(new ContentId[] {PIET_SHARED_STATE_2}, secondScope, /* shouldExist= */ true);
   }
 
   @Test
   public void testGc_contentInExpiredSessionDeleted() {
     scope
         .getFakeFeedRequestManager()
-        .queueResponse(createResponse(REQUEST_1))
+        .queueResponse(createResponse(REQUEST_1, PIET_SHARED_STATE_1))
         .triggerRefresh(
             RequestReason.OPEN_WITHOUT_CONTENT,
             scope.getFeedSessionManager().getUpdateConsumer(MutationContext.EMPTY_CONTEXT));
@@ -104,7 +120,7 @@
     // Populate HEAD with new data.
     scope
         .getFakeFeedRequestManager()
-        .queueResponse(createResponse(REQUEST_2))
+        .queueResponse(createResponse(REQUEST_2, PIET_SHARED_STATE_2))
         .triggerRefresh(
             RequestReason.OPEN_WITHOUT_CONTENT,
             scope.getFeedSessionManager().getUpdateConsumer(MutationContext.EMPTY_CONTEXT));
@@ -114,7 +130,9 @@
     fakeClock.advance(LIFETIME_MS + 1L);
     InfraIntegrationScope secondScope = scope.clone();
     assertPayloads(REQUEST_1, secondScope, /* shouldExist= */ false);
+    assertSharedStates(new ContentId[] {PIET_SHARED_STATE_1}, secondScope, /* shouldExist= */ true);
     assertPayloads(REQUEST_2, secondScope, /* shouldExist= */ true);
+    assertSharedStates(new ContentId[] {PIET_SHARED_STATE_2}, secondScope, /* shouldExist= */ true);
   }
 
   private static void assertPayloads(
@@ -136,7 +154,31 @@
     }
   }
 
-  private static Response createResponse(ContentId[] contentIds) {
-    return ResponseBuilder.forClearAllWithCards(contentIds).build();
+  private static void assertSharedStates(
+      ContentId[] contentIds, InfraIntegrationScope scope, boolean shouldExist) {
+    scope.getFakeThreadUtils().enforceMainThread(false);
+    List<String> sharedStateContentIds = new ArrayList<>();
+    for (StreamSharedState streamSharedState : scope.getStore().getSharedStates().getValue()) {
+      sharedStateContentIds.add(streamSharedState.getContentId());
+    }
+    List<String> expectedContentIds = new ArrayList<>(contentIds.length);
+    for (ContentId contentId : contentIds) {
+      expectedContentIds.add(scope.getProtocolAdapter().getStreamContentId(contentId));
+    }
+
+    if (shouldExist) {
+      assertThat(sharedStateContentIds).containsAtLeastElementsIn(expectedContentIds);
+    } else {
+      assertThat(sharedStateContentIds).containsNoneIn(expectedContentIds);
+    }
+  }
+
+  private static Response createResponse(ContentId[] cards, ContentId pietSharedStateContentId) {
+    return ResponseBuilder.builder()
+        .addClearOperation()
+        .addPietSharedState(pietSharedStateContentId)
+        .addRootFeature()
+        .addCardsToRoot(cards)
+        .build();
   }
 }