make sure to check if the logDataBindingRef is present if logData is not

PiperOrigin-RevId: 253821396
Change-Id: I9d1b9f500425d755cd086f5181c511a294037409
diff --git a/src/main/java/com/google/android/libraries/feed/piet/ElementAdapter.java b/src/main/java/com/google/android/libraries/feed/piet/ElementAdapter.java
index abd2990..8017458 100644
--- a/src/main/java/com/google/android/libraries/feed/piet/ElementAdapter.java
+++ b/src/main/java/com/google/android/libraries/feed/piet/ElementAdapter.java
@@ -38,6 +38,7 @@
 import com.google.search.now.ui.piet.ElementsProto.Visibility;
 import com.google.search.now.ui.piet.ElementsProto.VisibilityState;
 import com.google.search.now.ui.piet.ErrorsProto.ErrorCode;
+import com.google.search.now.ui.piet.LogDataProto.LogData;
 import com.google.search.now.ui.piet.StylesProto.StyleIdsStack;
 import java.util.HashSet;
 import java.util.Set;
@@ -60,6 +61,7 @@
   private Element baseElement = Element.getDefaultInstance();
   /*@Nullable*/ private RecyclerKey key;
   private StyleProvider elementStyle;
+  /*@Nullable*/ private LogData logData = null;
 
   /**
    * Set to {@code true} when {@link #createAdapter} has completed successfully; unset at the end of
@@ -294,8 +296,21 @@
         .setContentDescription(
             getAccessibilityString(baseElement.getAccessibility(), frameContext));
     LogDataCallback callback = parameters.hostProviders.getLogDataCallback();
-    if (callback != null && baseElement.hasLogData()) {
-      callback.onBind(baseElement.getLogData(), getView());
+    if (callback != null) {
+      switch (baseElement.getLogDataValueCase()) {
+        case LOG_DATA:
+          logData = baseElement.getLogData();
+          break;
+        case LOG_DATA_REF:
+          // Get the logData from the binding in the template
+          logData = frameContext.getLogDataFromBinding(baseElement.getLogDataRef());
+          break;
+        default: // No log Data
+      }
+      if (logData != null) {
+
+        callback.onBind(logData, getView());
+      }
     }
   }
 
@@ -421,8 +436,9 @@
       onUnbindModel();
     }
     LogDataCallback callback = parameters.hostProviders.getLogDataCallback();
-    if (callback != null && baseElement.hasLogData()) {
-      callback.onUnbind(baseElement.getLogData(), getView());
+    if (callback != null && logData != null) {
+      callback.onUnbind(logData, getView());
+      logData = null;
     }
     model = null;
     baseElement = Element.getDefaultInstance();
diff --git a/src/main/java/com/google/android/libraries/feed/piet/FrameContext.java b/src/main/java/com/google/android/libraries/feed/piet/FrameContext.java
index 232e592..b326121 100644
--- a/src/main/java/com/google/android/libraries/feed/piet/FrameContext.java
+++ b/src/main/java/com/google/android/libraries/feed/piet/FrameContext.java
@@ -36,6 +36,7 @@
 import com.google.search.now.ui.piet.BindingRefsProto.ElementBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.GridCellWidthBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ImageBindingRef;
+import com.google.search.now.ui.piet.BindingRefsProto.LogDataBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ParameterizedTextBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.StyleBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.TemplateBindingRef;
@@ -47,6 +48,7 @@
 import com.google.search.now.ui.piet.ErrorsProto.ErrorCode;
 import com.google.search.now.ui.piet.ImagesProto.Image;
 import com.google.search.now.ui.piet.ImagesProto.ImageSource;
+import com.google.search.now.ui.piet.LogDataProto.LogData;
 import com.google.search.now.ui.piet.PietProto.Frame;
 import com.google.search.now.ui.piet.PietProto.PietSharedState;
 import com.google.search.now.ui.piet.PietProto.Stylesheet;
@@ -552,6 +554,31 @@
     }
   }
 
+  /**
+   * Return an {@link LogData} for the binding if there is one defined; otherwise returns the
+   * default instance.
+   */
+  /*@Nullable*/
+  LogData getLogDataFromBinding(LogDataBindingRef binding) {
+    BindingValue bindingValue = bindingValues.get(binding.getBindingId());
+    if (bindingValue != null && bindingValue.hasHostBindingData()) {
+      bindingValue = hostBindingProvider.getLogDataBindingForValue(bindingValue);
+    }
+    if (bindingValue == null) {
+      return null;
+    }
+    if (!bindingValue.hasLogData()) {
+      Logger.w(
+          TAG,
+          reportMessage(
+              MessageType.WARNING,
+              ERR_BINDING_VALUE_TYPE_MISMATCH,
+              String.format("No logData found for binding %s", binding.getBindingId())));
+      return null;
+    }
+    return bindingValue.getLogData();
+  }
+
   /** Returns the root view of this Frame. */
   View getFrameView() {
     return frameView;
diff --git a/src/test/java/com/google/android/libraries/feed/piet/ElementAdapterTest.java b/src/test/java/com/google/android/libraries/feed/piet/ElementAdapterTest.java
index 43f96b3..0a87520 100644
--- a/src/test/java/com/google/android/libraries/feed/piet/ElementAdapterTest.java
+++ b/src/test/java/com/google/android/libraries/feed/piet/ElementAdapterTest.java
@@ -53,6 +53,7 @@
 import com.google.search.now.ui.piet.ActionsProto.VisibilityAction;
 import com.google.search.now.ui.piet.BindingRefsProto.ActionsBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ElementBindingRef;
+import com.google.search.now.ui.piet.BindingRefsProto.LogDataBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ParameterizedTextBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.VisibilityBindingRef;
 import com.google.search.now.ui.piet.ElementsProto.BindingValue;
@@ -88,7 +89,7 @@
   private static final boolean LEGACY_CORNERS_FLAG = false;
   private static final boolean OUTLINE_CORNERS_FLAG = false;
 
-  private final LogData logDataTest = LogData.getDefaultInstance();
+  private final LogData logDataTest = LogData.newBuilder().build();
 
   @Mock private FrameContext frameContext;
   @Mock private HostProviders hostProviders;
@@ -519,6 +520,17 @@
   }
 
   @Test
+  public void testBindModel_callsOnBindLogDataCallback_bindingRef() {
+    LogDataBindingRef logDataBinding =
+        LogDataBindingRef.newBuilder().setBindingId("LogData!").build();
+    Element defaultElement = Element.newBuilder().setLogDataRef(logDataBinding).build();
+    when(frameContext.getLogDataFromBinding(logDataBinding)).thenReturn(logDataTest);
+    adapter.bindModel(defaultElement, defaultElement, frameContext);
+
+    assertThat(callbackBound).isTrue();
+  }
+
+  @Test
   public void bindModel_extractsModelFromElement() {
     Element element = Element.getDefaultInstance();
 
@@ -869,6 +881,18 @@
   }
 
   @Test
+  public void testUnbindModel_callsOnBindLogDataCallback_bindingRef() {
+    LogDataBindingRef logDataBinding =
+        LogDataBindingRef.newBuilder().setBindingId("LogData!").build();
+    Element defaultElement = Element.newBuilder().setLogDataRef(logDataBinding).build();
+    when(frameContext.getLogDataFromBinding(logDataBinding)).thenReturn(logDataTest);
+
+    adapter.bindModel(defaultElement, defaultElement, frameContext);
+    adapter.unbindModel();
+    assertThat(callbackUnbound).isTrue();
+  }
+
+  @Test
   public void unbindModel_unsetsModel() {
     Element element = Element.getDefaultInstance();
 
diff --git a/src/test/java/com/google/android/libraries/feed/piet/FrameContextTest.java b/src/test/java/com/google/android/libraries/feed/piet/FrameContextTest.java
index b202b65..9de507d 100644
--- a/src/test/java/com/google/android/libraries/feed/piet/FrameContextTest.java
+++ b/src/test/java/com/google/android/libraries/feed/piet/FrameContextTest.java
@@ -39,6 +39,7 @@
 import com.google.search.now.ui.piet.BindingRefsProto.ElementBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.GridCellWidthBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ImageBindingRef;
+import com.google.search.now.ui.piet.BindingRefsProto.LogDataBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.ParameterizedTextBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.StyleBindingRef;
 import com.google.search.now.ui.piet.BindingRefsProto.TemplateBindingRef;
@@ -58,6 +59,7 @@
 import com.google.search.now.ui.piet.GradientsProto.LinearGradient;
 import com.google.search.now.ui.piet.ImagesProto.Image;
 import com.google.search.now.ui.piet.ImagesProto.ImageSource;
+import com.google.search.now.ui.piet.LogDataProto.LogData;
 import com.google.search.now.ui.piet.MediaQueriesProto.ComparisonCondition;
 import com.google.search.now.ui.piet.MediaQueriesProto.DarkLightCondition;
 import com.google.search.now.ui.piet.MediaQueriesProto.DarkLightCondition.DarkLightMode;
@@ -714,6 +716,41 @@
   }
 
   @Test
+  public void testGetLogDataFromBinding() {
+    LogData logData = LogData.newBuilder().build();
+    frameContext = makeFrameContextWithBinding(defaultBinding().setLogData(logData).build());
+    assertThat(
+            frameContext.getLogDataFromBinding(
+                LogDataBindingRef.newBuilder().setBindingId(BINDING_ID).build()))
+        .isSameInstanceAs(logData);
+
+    assertThat(
+            frameContext.getLogDataFromBinding(
+                LogDataBindingRef.newBuilder().setBindingId(INVALID_BINDING_ID).build()))
+        .isNull();
+
+    frameContext = makeFrameContextWithNoBindings();
+    assertThat(
+            frameContext.getLogDataFromBinding(
+                LogDataBindingRef.newBuilder().setBindingId(BINDING_ID).build()))
+        .isNull();
+  }
+
+  @Test
+  public void testGetLogDataFromBinding_hostBinding() {
+    frameContext =
+        makeFrameContextWithBinding(
+            defaultBinding()
+                .setHostBindingData(HostBindingData.newBuilder())
+                .setLogData(LogData.getDefaultInstance())
+                .build());
+    assertThat(
+            frameContext.getLogDataFromBinding(
+                LogDataBindingRef.newBuilder().setBindingId(BINDING_ID).build()))
+        .isEqualTo(LogData.getDefaultInstance());
+  }
+
+  @Test
   public void testGetStyleFromBinding() {
     BoundStyle boundStyle = BoundStyle.newBuilder().setColor(12345).build();
     frameContext = makeFrameContextWithBinding(defaultBinding().setBoundStyle(boundStyle).build());