Clean up proxy after unit tests (#932)

snowflakedb · Jan 29, 2025 · 7193c5d · 7193c5d
1 parent 6c3190c
commit 7193c5d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 19 deletions.
diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/InternalStageTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/InternalStageTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2021-2025 Snowflake Computing Inc. All rights reserved.
  */
 
 package net.snowflake.ingest.streaming.internal;
@@ -539,18 +539,12 @@ public void testGenerateProxyPropertiesForJDBC() {
           nonProxyHosts, props.get(SFSessionProperty.NON_PROXY_HOSTS.getPropertyKey()));
     } finally {
       // Cleanup
-      if (oldUseProxy != null) {
-        System.setProperty(USE_PROXY, oldUseProxy);
-        System.setProperty(PROXY_HOST, oldProxyHost);
-        System.setProperty(PROXY_PORT, oldProxyPort);
-      }
-      if (oldUser != null) {
-        System.setProperty(HTTP_PROXY_USER, oldUser);
-        System.setProperty(HTTP_PROXY_PASSWORD, oldPassword);
-      }
-      if (oldNonProxyHosts != null) {
-        System.setProperty(NON_PROXY_HOSTS, oldNonProxyHosts);
-      }
+      resetProperty(USE_PROXY, oldUseProxy);
+      resetProperty(PROXY_HOST, oldProxyHost);
+      resetProperty(PROXY_PORT, oldProxyPort);
+      resetProperty(HTTP_PROXY_USER, oldUser);
+      resetProperty(HTTP_PROXY_PASSWORD, oldPassword);
+      resetProperty(NON_PROXY_HOSTS, oldNonProxyHosts);
     }
   }
 
@@ -584,9 +578,7 @@ public void testShouldBypassProxy() {
     Assert.assertFalse(shouldBypassProxy(accountUnderscoreName));
     Assert.assertFalse(shouldBypassProxy(accountNamePrivateLink));
 
-    if (oldNonProxyHosts != null) {
-      System.setProperty(NON_PROXY_HOSTS, oldNonProxyHosts);
-    }
+    resetProperty(NON_PROXY_HOSTS, oldNonProxyHosts);
   }
 
   @Test
@@ -663,4 +655,12 @@ private HttpEntity createHttpEntity(String content) {
     entity.setContent(new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8)));
     return entity;
   }
+
+  private void resetProperty(String key, String oldValue) {
+    if (oldValue != null) {
+      System.setProperty(key, oldValue);
+    } else {
+      System.clearProperty(key);
+    }
+  }
 }
diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestIT.java
@@ -1,5 +1,10 @@
+/*
+ * Copyright (c) 2025 Snowflake Computing Inc. All rights reserved.
+ */
+
 package net.snowflake.ingest.streaming.internal;
 
+import static net.snowflake.ingest.TestUtils.waitForOffset;
 import static net.snowflake.ingest.utils.Constants.BLOB_NO_HEADER;
 import static net.snowflake.ingest.utils.Constants.COMPRESS_BLOB_TWICE;
 import static net.snowflake.ingest.utils.Constants.DROP_CHANNEL_ENDPOINT;
@@ -75,7 +80,7 @@ public class StreamingIngestIT {
   private Connection jdbcConnection;
   private String testDb;
 
-  @Parameters
+  @Parameters(name = "compressionAlgorithm={0}, enableIcebergStreaming={1}")
   public static Iterable<Object[]> getParameterPermutations() {
     return Arrays.asList(
         new Object[][] {{"GZIP", false}, {"ZSTD", false}, {"GZIP", true}, {"ZSTD", true}});
@@ -838,12 +843,14 @@ public void testTwoClientsOneChannel() throws Exception {
     row.put("c1", "1");
     verifyInsertValidationResponse(channelA.insertRow(row, "1"));
     clientA.flush(false).get();
+    waitForOffset(channelA, "1");
 
     // ClientB opens channel and invalidates ClientA
     SnowflakeStreamingIngestChannel channelB = clientB.openChannel(requestA);
     row.put("c1", "2");
     verifyInsertValidationResponse(channelB.insertRow(row, "2"));
     clientB.flush(false).get();
+    waitForOffset(channelB, "2");
 
     // ClientA tries to write, but will fail to register because it's invalid
     row.put("c1", "3");
@@ -1270,7 +1277,7 @@ public void testFailureHalfwayThroughRowProcessing() throws Exception {
     row.put("c2", null);
     channel.insertRow(row, "offset2");
 
-    TestUtils.waitForOffset(channel, "offset2");
+    waitForOffset(channel, "offset2");
 
     jdbcConnection.createStatement().execute(String.format("alter table %s migrate;", tableName));
 
@@ -1462,7 +1469,7 @@ private void ingestByColIndexAndVerifyTable(
       }
       verifyInsertValidationResponse(channel.insertRow(row, Integer.toString(i)));
     }
-    TestUtils.waitForOffset(channel, Integer.toString(numberOfRows - 1));
+    waitForOffset(channel, Integer.toString(numberOfRows - 1));
     channel.close();
 
     // query them and verify