From 4ca5acdc33812ee863d1613ab7e55001c8d12cf0 Mon Sep 17 00:00:00 2001
From: Krishna Anandan <krishna1729atom@gmail.com>
Date: Wed, 15 Nov 2023 07:10:09 -0600
Subject: [PATCH] Fixed 2 Flaky Tests (#15376)

---
 .../data/input/protobuf/ProtobufInputFormatTest.java   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/extensions-core/protobuf-extensions/src/test/java/org/apache/druid/data/input/protobuf/ProtobufInputFormatTest.java b/extensions-core/protobuf-extensions/src/test/java/org/apache/druid/data/input/protobuf/ProtobufInputFormatTest.java
index 56b774f15ee..8fd280ca7e4 100644
--- a/extensions-core/protobuf-extensions/src/test/java/org/apache/druid/data/input/protobuf/ProtobufInputFormatTest.java
+++ b/extensions-core/protobuf-extensions/src/test/java/org/apache/druid/data/input/protobuf/ProtobufInputFormatTest.java
@@ -24,6 +24,7 @@ import com.fasterxml.jackson.databind.Module;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.io.Files;
 import org.apache.druid.common.config.NullHandling;
@@ -57,6 +58,7 @@ import org.junit.rules.ExpectedException;
 import java.io.File;
 import java.io.IOException;
 import java.util.Collections;
+import java.util.HashSet;
 
 public class ProtobufInputFormatTest
 {
@@ -231,7 +233,7 @@ public class ProtobufInputFormatTest
     ).read().next();
 
     Assert.assertEquals(
-        ImmutableList.builder()
+        ImmutableSet.builder()
                      .add("eventType")
                      .add("foobar")
                      .add("bar0")
@@ -244,7 +246,7 @@ public class ProtobufInputFormatTest
                      .add("id")
                      .add("someBytesColumn")
                      .build(),
-        row.getDimensions()
+        new HashSet<>(row.getDimensions())
     );
 
     ProtobufInputRowParserTest.verifyNestedData(row, dateTime);
@@ -368,7 +370,7 @@ public class ProtobufInputFormatTest
     InputRow row = transformingReader.read().next();
 
     Assert.assertEquals(
-        ImmutableList.of(
+        ImmutableSet.of(
             "someOtherId",
             "someIntColumn",
             "isValid",
@@ -381,7 +383,7 @@ public class ProtobufInputFormatTest
             "id",
             "someBytesColumn"
         ),
-        row.getDimensions()
+        new HashSet<>(row.getDimensions())
     );
 
     Assert.assertEquals(ImmutableMap.of("bar", "baz"), row.getRaw("foo"));