From 6dcad1e4905058ff45eb339edce0dbd892447cbc Mon Sep 17 00:00:00 2001
From: Purujit Saha <purujit.saha@snowflake.com>
Date: Thu, 31 Oct 2024 18:03:56 +0000
Subject: [PATCH] Fix test to handle iceberg case

---
 .../streaming/internal/RowBufferTest.java     | 40 ++++++++++++-------
 1 file changed, 25 insertions(+), 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
index 5402aba95..747773e9e 100644
--- 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;
+        }
+      }
     }
   }