Fix types

apache · Dec 6, 2016 · 911cfd8 · 911cfd8
1 parent c76f045
commit 911cfd8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/display/DisplayData.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/display/DisplayData.java
@@ -866,9 +866,16 @@ public static ItemSpec<String> item(String key, @Nullable String value) {
   /**
    * Create a display item for the specified key and {@link ValueProvider}.
    */
-  public static ItemSpec<String> item(String key, ValueProvider<?> value) {
-    return item(key, Type.STRING, value.isAccessible()
-        ? value.get().toString() : value.toString());
+  public static ItemSpec item(String key, ValueProvider<?> value) {
+    if (value.isAccessible()) {
+      Object got = value.get();
+      Type type = inferType(got);
+      if (type == null) {
+        throw new RuntimeException(String.format("Unknown value type: %s", got));
+      }
+      return item(key, type, got);
+    }
+    return item(key, Type.STRING, value.toString());
   }
 
   /**

diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/display/DisplayDataTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/display/DisplayDataTest.java
@@ -166,6 +166,30 @@ public void populateDisplayData(DisplayData.Builder builder) {
     assertThat(data, hasDisplayItem("foo", "bar"));
   }
 
+  @Test
+  public void testStaticValueProviderDate() {
+    final Instant value = Instant.now();
+    DisplayData data =
+        DisplayData.from(new HasDisplayData() {
+              @Override
+              public void populateDisplayData(DisplayData.Builder builder) {
+                builder.add(DisplayData.item(
+                    "foo", StaticValueProvider.of(value)));
+              }
+            });
+
+    @SuppressWarnings("unchecked")
+    DisplayData.Item item = (DisplayData.Item) data.items().toArray()[0];
+
+    @SuppressWarnings("unchecked")
+    Matcher<Item> matchesAllOf = Matchers.allOf(
+        hasKey("foo"),
+        hasType(DisplayData.Type.TIMESTAMP),
+        hasValue(ISO_FORMATTER.print(value)));
+
+    assertThat(item, matchesAllOf);
+  }
+
   @Test
   public void testStaticValueProviderString() {
     DisplayData data =
@@ -193,7 +217,7 @@ public void populateDisplayData(DisplayData.Builder builder) {
             });
 
     assertThat(data.items(), hasSize(1));
-    assertThat(data, hasDisplayItem("foo", "1"));
+    assertThat(data, hasDisplayItem("foo", 1));
   }
 
   @Test