fix test

snowflakedb · Jun 11, 2024 · 86a7934 · 86a7934
1 parent 9dbc608
commit 86a7934
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java
@@ -152,6 +152,9 @@ List<List<ChannelData<T>>> getData() {
     this.owningClient = client;
     this.channelCache = cache;
     this.targetStage = targetStage;
+    if (this.targetStage != null) {
+      this.clientPrefix = this.targetStage.getClientPrefix();
+    }
     this.counter = new AtomicLong(0);
     this.registerService = new RegisterService<>(client, isTestMode);
     this.isNeedFlush = false;
@@ -188,7 +191,7 @@ List<List<ChannelData<T>>> getData() {
                 client.getRequestBuilder(),
                 client.getName(),
                 DEFAULT_MAX_UPLOAD_RETRIES);
-        this.clientPrefix = this.targetStage.clientPrefix;
+        this.clientPrefix = this.targetStage.getClientPrefix();
       } else {
         this.targetStage = null;
       }
@@ -730,7 +733,7 @@ private String getBlobPath(String volumeHash) {
 
   /** For TESTING */
   String getBlobPath(Calendar calendar, String volumeHash) {
-    if (isTestMode) {
+    if (isTestMode && this.clientPrefix == null) {
       this.clientPrefix = "testPrefix";
     }
     Utils.assertStringNotNullOrEmpty("client prefix", this.clientPrefix);

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java
@@ -408,10 +408,9 @@ public void testGetFilePath() {
     TestContext<?> testContext = testContextFactory.create(false);
     FlushService<?> flushService = testContext.flushService;
     Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC"));
-    String clientPrefix = "honk";
     String outputString = flushService.getBlobPath(calendar, null);
     Path outputPath = Paths.get(outputString);
-    Assert.assertTrue(outputPath.getFileName().toString().contains(clientPrefix));
+    Assert.assertTrue(outputPath.getFileName().toString().contains("client_prefix"));
     Assert.assertTrue(
         calendar.get(Calendar.MINUTE)
                 - Integer.parseInt(outputPath.getParent().getFileName().toString())