Set bounds on the wrapper LayerDrawable in ChunkedTextElementAdapter after the image loads.

This will let us adjust the size of the image if we don't know the size beforehand.

PiperOrigin-RevId: 277759348
Change-Id: I6ac59d333483b6738324567fd77520d415a07902
diff --git a/src/main/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapter.java b/src/main/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapter.java
index d75c352..8e6f1c5 100644
--- a/src/main/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapter.java
+++ b/src/main/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapter.java
@@ -493,6 +493,7 @@
           wrapper.setDrawableByLayerId(SINGLE_LAYER_ID, finalDrawable),
           "Failed to set drawable on chunked text");
       setBounds(finalDrawable, imageStyle, textView);
+      setBounds(wrapper, imageStyle, textView);
       textView.invalidate();
       loadingImages.remove(this);
     }
diff --git a/src/test/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapterTest.java b/src/test/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapterTest.java
index b58322be..a00de72 100644
--- a/src/test/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapterTest.java
+++ b/src/test/java/com/google/android/libraries/feed/piet/ChunkedTextElementAdapterTest.java
@@ -624,6 +624,58 @@
 
     assertThat(imageDrawable.getBounds())
         .isEqualTo(new Rect(0, 0, STYLE_WIDTH_PX, STYLE_HEIGHT_PX));
+    assertThat(containerDrawable.getBounds())
+        .isEqualTo(new Rect(0, 0, STYLE_WIDTH_PX, STYLE_HEIGHT_PX));
+  }
+
+  @Test
+  public void testAddImageChunk_setsImageAndDimsWidthNotSpecified() {
+    when(mockStyleProvider.hasWidth()).thenReturn(true);
+    when(mockStyleProvider.getWidthSpecPx(context)).thenReturn(-1);
+    when(mockStyleProvider.hasHeight()).thenReturn(true);
+    when(mockStyleProvider.getHeightSpecPx(context)).thenReturn(STYLE_HEIGHT_PX);
+    when(mockStyleProvider.hasPreLoadFill()).thenReturn(true);
+    Drawable preLoadFill = new ColorDrawable(Color.CYAN);
+    when(mockStyleProvider.createPreLoadFill()).thenReturn(preLoadFill);
+
+    // Required to set up the local frameContext member var.
+    adapter.createAdapter(asElement(TextElement.getDefaultInstance()), frameContext);
+
+    adapter.addImageChunk(frameContext, textView, spannable, IMAGE_CHUNK);
+
+    ImageSpan[] imageSpans = spannable.getSpans(0, 1, ImageSpan.class);
+    LayerDrawable containerDrawable = (LayerDrawable) imageSpans[0].getDrawable();
+
+    ArgumentCaptor<ImageSpanDrawableCallback> imageCallbackCaptor =
+        ArgumentCaptor.forClass(ImageSpanDrawableCallback.class);
+    verify(mockAssetProvider)
+        .getImage(
+            eq(IMAGE_CHUNK_IMAGE), eq(-1), eq(STYLE_HEIGHT_PX), imageCallbackCaptor.capture());
+
+    int aspectRatio = 4;
+    Drawable imageDrawable =
+        new ColorDrawable(123) {
+          @Override
+          public int getIntrinsicWidth() {
+            return 11 * aspectRatio;
+          }
+
+          @Override
+          public int getIntrinsicHeight() {
+            return 11;
+          }
+        };
+
+    // Activate the image loading callback
+    imageCallbackCaptor.getValue().accept(imageDrawable);
+
+    // Assert that we set the image on the span
+    assertThat(containerDrawable.getDrawable(0)).isSameInstanceAs(imageDrawable);
+
+    assertThat(imageDrawable.getBounds())
+        .isEqualTo(new Rect(0, 0, STYLE_HEIGHT_PX * 4, STYLE_HEIGHT_PX));
+    assertThat(containerDrawable.getBounds())
+        .isEqualTo(new Rect(0, 0, STYLE_HEIGHT_PX * 4, STYLE_HEIGHT_PX));
   }
 
   @Test