Fix test to handle iceberg case

snowflakedb · Oct 31, 2024 · 6dcad1e · 6dcad1e
1 parent 4008e23
commit 6dcad1e
Showing 1 changed file with 25 additions and 15 deletions.
diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java
@@ -2047,21 +2047,31 @@ public void testParquetFileNameMetadata() throws IOException {
           reader.getKeyValueMetadata().get(Constants.SDK_VERSION_KEY));
     }
     {
-      Flusher.SerializationResult result =
-          flusher.serialize(Collections.singletonList(data), filePath, 13);
-
-      BdecParquetReader reader = new BdecParquetReader(result.chunkData.toByteArray());
-      Assert.assertEquals(
-          "testParquetFileNameMetadata_13.bdec",
-          reader
-              .getKeyValueMetadata()
-              .get(
-                  enableIcebergStreaming
-                      ? Constants.ASSIGNED_FULL_FILE_NAME_KEY
-                      : Constants.PRIMARY_FILE_ID_KEY));
-      Assert.assertEquals(
-          RequestBuilder.DEFAULT_VERSION,
-          reader.getKeyValueMetadata().get(Constants.SDK_VERSION_KEY));
+      try {
+        Flusher.SerializationResult result =
+            flusher.serialize(Collections.singletonList(data), filePath, 13);
+        if (enableIcebergStreaming) {
+          Assert.fail(
+              "Should have thrown an exception because iceberg streams do not support offsets");
+        }
+
+        BdecParquetReader reader = new BdecParquetReader(result.chunkData.toByteArray());
+        Assert.assertEquals(
+            "testParquetFileNameMetadata_13.bdec",
+            reader
+                .getKeyValueMetadata()
+                .get(
+                    enableIcebergStreaming
+                        ? Constants.ASSIGNED_FULL_FILE_NAME_KEY
+                        : Constants.PRIMARY_FILE_ID_KEY));
+        Assert.assertEquals(
+            RequestBuilder.DEFAULT_VERSION,
+            reader.getKeyValueMetadata().get(Constants.SDK_VERSION_KEY));
+      } catch (IllegalStateException ex) {
+        if (!enableIcebergStreaming) {
+          throw ex;
+        }
+      }
     }
   }