Add tests for SnowflakeFileTransferAgent (part 1) (#1288)

* add tests for SnowflakeFileTransferAgent * add isInjectedFileTransferExceptionEnabled()
snowflakedb · Mar 13, 2023 · 7cd122d · 7cd122d
1 parent f7ae088
commit 7cd122d
Show file tree

Hide file tree

Showing 5 changed files with 369 additions and 7 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java
@@ -96,9 +96,18 @@ public class SnowflakeFileTransferAgent extends SFBaseFileTransferAgent {
 
   // default parallelism
   private int parallel = DEFAULT_PARALLEL;
-
   private SFSession session;
   private SFStatement statement;
+  private static Throwable injectedFileTransferException = null; // for testing purpose
+
+  // This function should only be used for testing purpose
+  static void setInjectedFileTransferException(Throwable th) {
+    injectedFileTransferException = th;
+  }
+
+  static boolean isInjectedFileTransferExceptionEnabled() {
+    return injectedFileTransferException != null;
+  }
 
   public StageInfo getStageInfo() {
     return this.stageInfo;
@@ -315,6 +324,13 @@ private static InputStreamWithMetadata compressStreamWithGZIP(
 
       countingStream.flush();
 
+      // Normal flow will never hit here. This is only for testing purposes
+      if (isInjectedFileTransferExceptionEnabled()
+          && SnowflakeFileTransferAgent.injectedFileTransferException
+              instanceof NoSuchAlgorithmException) {
+        throw (NoSuchAlgorithmException) SnowflakeFileTransferAgent.injectedFileTransferException;
+      }
+
       return new InputStreamWithMetadata(
           countingStream.getCount(),
           Base64.getEncoder().encodeToString(digestStream.getMessageDigest().digest()),
@@ -362,6 +378,10 @@ private static InputStreamWithMetadata compressStreamWithGZIPNoDigest(
 
       countingStream.flush();
 
+      // Normal flow will never hit here. This is only for testing purposes
+      if (isInjectedFileTransferExceptionEnabled()) {
+        throw (IOException) SnowflakeFileTransferAgent.injectedFileTransferException;
+      }
       return new InputStreamWithMetadata(countingStream.getCount(), null, tempStream);
 
     } catch (IOException ex) {
@@ -765,6 +785,12 @@ private void parseCommand() throws SnowflakeSQLException {
     String[] src_locations;
 
     try {
+      // Normal flow will never hit here. This is only for testing purposes
+      if (isInjectedFileTransferExceptionEnabled()
+          && injectedFileTransferException instanceof SnowflakeSQLException) {
+        throw (SnowflakeSQLException) SnowflakeFileTransferAgent.injectedFileTransferException;
+      }
+
       src_locations = mapper.readValue(locationsNode.toString(), String[].class);
       initEncryptionMaterial(commandType, jsonNode);
       initPresignedUrls(commandType, jsonNode);
@@ -1696,6 +1722,12 @@ static Set<String> expandFileNames(String[] filePathList) throws SnowflakeSQLExc
             entry.getKey(),
             entry.getValue().toString());
 
+        // Normal flow will never hit here. This is only for testing purposes
+        if (isInjectedFileTransferExceptionEnabled()
+            && injectedFileTransferException instanceof Exception) {
+          throw (Exception) SnowflakeFileTransferAgent.injectedFileTransferException;
+        }
+
         // The following currently ignore sub directories
         for (Object file :
             FileUtils.listFiles(dir, new WildcardFileFilter(entry.getValue()), null)) {
@@ -1963,6 +1995,11 @@ public static void uploadWithoutConnection(SnowflakeFileTransferConfig config) t
       SnowflakeStorageClient initialClient =
           storageFactory.createClient(stageInfo, 1, encMat, /*session = */ null);
 
+      // Normal flow will never hit here. This is only for testing purposes
+      if (isInjectedFileTransferExceptionEnabled()) {
+        throw (Exception) SnowflakeFileTransferAgent.injectedFileTransferException;
+      }
+
       switch (stageInfo.getStageType()) {
         case S3:
         case AZURE:
@@ -2438,6 +2475,11 @@ private void compareAndSkipRemoteFiles(
       final boolean remoteEncrypted;
 
       try {
+        // Normal flow will never hit here. This is only for testing purposes
+        if (isInjectedFileTransferExceptionEnabled()) {
+          throw (NoSuchAlgorithmException) SnowflakeFileTransferAgent.injectedFileTransferException;
+        }
+
         localFile =
             (commandType == CommandType.UPLOAD) ? mappedSrcFile : (localLocation + objFileName);
 

diff --git a/src/test/java/net/snowflake/client/jdbc/FileUploaderExpandFileNamesTest.java b/src/test/java/net/snowflake/client/jdbc/FileUploaderExpandFileNamesTest.java
@@ -10,6 +10,7 @@
 import java.util.Properties;
 import java.util.Set;
 import net.snowflake.client.core.OCSPMode;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -39,6 +40,21 @@ public void testProcessFileNames() throws Exception {
     assertTrue(files.contains(folderName + "/TestFileD"));
   }
 
+  @Test
+  public void testProcessFileNamesException() {
+    // inject the Exception
+    SnowflakeFileTransferAgent.setInjectedFileTransferException(new Exception());
+    String[] locations = {"/Tes*Fil*A", "/TestFil?B", "~/TestFileC", "TestFileD"};
+
+    try {
+      SnowflakeFileTransferAgent.expandFileNames(locations);
+    } catch (SnowflakeSQLException err) {
+      Assert.assertEquals(200007, err.getErrorCode());
+      Assert.assertEquals("22000", err.getSQLState());
+    }
+    SnowflakeFileTransferAgent.setInjectedFileTransferException(null);
+  }
+
   @Test
   public void testSnowflakeFileTransferConfig() throws Exception {
     SnowflakeFileTransferMetadataV1 metadata =