Determine pubsub event timestamp correctly in retry-timeout. (GoogleCloudPlatform#3770)

eamonnmcmanus · web-flow · commit 5e6957310960 · 2020-09-22T09:45:13.000-07:00
In the retry-timeout sample for GCF java11, we should determine the event timestamp from the Context parameter, not from properties of the passed-in PubSub object. Fixes GoogleCloudPlatform#3734.
diff --git a/functions/concepts/retry-timeout/src/main/java/functions/RetryTimeout.java b/functions/concepts/retry-timeout/src/main/java/functions/RetryTimeout.java
@@ -21,7 +21,6 @@
 import com.google.cloud.functions.BackgroundFunction;
 import com.google.cloud.functions.Context;
 import com.google.gson.Gson;
-import com.google.gson.JsonObject;
 import functions.eventpojos.PubSubMessage;
 import java.time.Duration;
 import java.time.ZoneOffset;
@@ -42,25 +41,19 @@ public class RetryTimeout implements BackgroundFunction<PubSubMessage> {
   @Override
   public void accept(PubSubMessage message, Context context) {
     ZonedDateTime utcNow = ZonedDateTime.now(ZoneOffset.UTC);
-    ZonedDateTime timestamp = utcNow;
+    ZonedDateTime timestamp = ZonedDateTime.parse(context.timestamp());
 
-    String data = message.getData();
-    JsonObject body = gson.fromJson(data, JsonObject.class);
-    if (body != null && body.has("timestamp")) {
-      String tz = body.get("timestamp").getAsString();
-      timestamp = ZonedDateTime.parse(tz);
-    }
     long eventAge = Duration.between(timestamp, utcNow).toMillis();
 
     // Ignore events that are too old
     if (eventAge > MAX_EVENT_AGE) {
-      logger.info(String.format("Dropping event %s.", data));
+      logger.info(String.format("Dropping event with timestamp %s.", timestamp));
       return;
     }
 
     // Process events that are recent enough
     // To retry this invocation, throw an exception here
-    logger.info(String.format("Processing event %s.", data));
+    logger.info(String.format("Processing event with timestamp %s.", timestamp));
   }
 }
 // [END functions_tips_infinite_retries]
diff --git a/functions/concepts/retry-timeout/src/test/java/functions/RetryTimeoutTest.java b/functions/concepts/retry-timeout/src/test/java/functions/RetryTimeoutTest.java
@@ -17,22 +17,23 @@
 package functions;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
+import com.google.cloud.functions.Context;
 import com.google.common.testing.TestLogHandler;
 import com.google.gson.Gson;
 import functions.eventpojos.PubSubMessage;
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
-import java.util.Map;
 import java.util.logging.Logger;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
-import org.mockito.MockitoAnnotations;
 
 @RunWith(JUnit4.class)
 public class RetryTimeoutTest {
@@ -52,8 +53,6 @@ public static void beforeClass() {
 
   @Before
   public void beforeTest() {
-    MockitoAnnotations.initMocks(this);
-
     LOG_HANDLER.clear();
   }
 
@@ -65,41 +64,29 @@ public void afterTest() {
 
   @Test
   public void retryTimeout_handlesRetryMsg() {
-    String timestampData = gson.toJson(Map.of(
-        "timestamp", ZonedDateTime.now(ZoneOffset.UTC).toString()));
+    ZonedDateTime timestamp = ZonedDateTime.now(ZoneOffset.UTC);
+    Context mockContext = mock(Context.class);
+    when(mockContext.timestamp()).thenReturn(timestamp.toString());
 
     PubSubMessage pubsubMessage = new PubSubMessage();
-    pubsubMessage.setData(timestampData);
 
-    new RetryTimeout().accept(pubsubMessage, null);
+    new RetryTimeout().accept(pubsubMessage, mockContext);
 
     String logMessage = LOG_HANDLER.getStoredLogRecords().get(0).getMessage();
-    assertThat(String.format("Processing event %s.", timestampData)).isEqualTo(logMessage);
+    assertThat(logMessage).contains("Processing event with timestamp " + timestamp);
   }
 
   @Test
   public void retryTimeout_handlesStopMsg() {
-    String timestamp = ZonedDateTime.ofInstant(Instant.ofEpochMilli(0), ZoneOffset.UTC).toString();
-    String timestampData = gson.toJson(Map.of("timestamp", timestamp));
-
-
-    PubSubMessage pubsubMessage = new PubSubMessage();
-    pubsubMessage.setData(timestampData);
-
-    new RetryTimeout().accept(pubsubMessage, null);
+    ZonedDateTime timestamp = ZonedDateTime.ofInstant(Instant.ofEpochMilli(0), ZoneOffset.UTC);
+    Context mockContext = mock(Context.class);
+    when(mockContext.timestamp()).thenReturn(timestamp.toString());
 
-    String logMessage = LOG_HANDLER.getStoredLogRecords().get(0).getMessage();
-    assertThat(String.format("Dropping event %s.", timestampData)).isEqualTo(logMessage);
-  }
-
-  @Test
-  public void retryTimeout_handlesEmptyMsg() {
     PubSubMessage pubsubMessage = new PubSubMessage();
-    pubsubMessage.setData("");
 
-    new RetryTimeout().accept(new PubSubMessage(), null);
+    new RetryTimeout().accept(pubsubMessage, mockContext);
 
     String logMessage = LOG_HANDLER.getStoredLogRecords().get(0).getMessage();
-    assertThat("Processing event null.").isEqualTo(logMessage);
+    assertThat(logMessage).contains("Dropping event with timestamp " + timestamp);
   }
 }
