From 4cf21b37aecb5845e69dc3b8d949ac5fbf5de522 Mon Sep 17 00:00:00 2001
From: Lukas Sembera <lukas.sembera@snowflake.com>
Date: Fri, 19 Jul 2024 14:33:37 +0000
Subject: [PATCH] Fix flaky test

---
 .../snowflake/ingest/streaming/internal/StreamingIngestIT.java | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java
index 2ecaa4e01..9426d4dc1 100644
--- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java
+++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java
@@ -8,6 +8,7 @@
 import static net.snowflake.ingest.utils.ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.core.Is.is;
+import static org.mockito.Mockito.atLeastOnce;
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
@@ -187,7 +188,7 @@ public void testSimpleIngest() throws Exception {
     // verify expected request sent to server
     String[] expectedPayloadParams = {"request_id", "blobs", "role", "blob_stats"};
     for (String expectedParam : expectedPayloadParams) {
-      Mockito.verify(requestBuilder)
+      Mockito.verify(requestBuilder, atLeastOnce())
           .generateStreamingIngestPostRequest(
               ArgumentMatchers.contains(expectedParam),
               ArgumentMatchers.refEq(REGISTER_BLOB_ENDPOINT),