Update tests for structured types

snowflakedb · Sep 2, 2024 · d5e0b1c · d5e0b1c
1 parent a0afebf
commit d5e0b1c
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 47 deletions.
diff --git a/src/test/java/com/snowflake/snowpark_test/JavaRowSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaRowSuite.java
@@ -571,29 +571,34 @@ public void getAsWithStructuredMap() {
   public void getAsWithStructuredArray() {
     structuredTypeTest(
         () -> {
-          TimeZone.setDefault(TimeZone.getTimeZone("US/Pacific"));
-
-          String query =
-              "SELECT "
-                  + "[1,2,3]::ARRAY(NUMBER) AS arr1,"
-                  + "['a','b']::ARRAY(VARCHAR) AS arr2,"
-                  + "[parse_json(31000000)::timestamp_ntz]::ARRAY(TIMESTAMP_NTZ) AS arr3,"
-                  + "[[1,2]]::ARRAY(ARRAY) AS arr4";
-
-          DataFrame df = getSession().sql(query);
-          Row row = df.collect()[0];
-
-          ArrayList<?> array1 = row.getAs(0, ArrayList.class);
-          assert array1.equals(Arrays.asList(1L, 2L, 3L));
-
-          ArrayList<?> array2 = row.getAs(1, ArrayList.class);
-          assert array2.equals(Arrays.asList("a", "b"));
-
-          ArrayList<?> array3 = row.getAs(2, ArrayList.class);
-          assert array3.equals(Collections.singletonList(new Timestamp(31000000000L)));
-
-          ArrayList<?> array4 = row.getAs(3, ArrayList.class);
-          assert array4.equals(Collections.singletonList("[\n  1,\n  2\n]"));
+          var oldTimeZone = TimeZone.getDefault();
+          try {
+            TimeZone.setDefault(TimeZone.getTimeZone("US/Pacific"));
+
+            String query =
+                "SELECT "
+                    + "[1,2,3]::ARRAY(NUMBER) AS arr1,"
+                    + "['a','b']::ARRAY(VARCHAR) AS arr2,"
+                    + "[parse_json(31000000)::timestamp_ntz]::ARRAY(TIMESTAMP_NTZ) AS arr3,"
+                    + "[[1,2]]::ARRAY(ARRAY) AS arr4";
+
+            DataFrame df = getSession().sql(query);
+            Row row = df.collect()[0];
+
+            ArrayList<?> array1 = row.getAs(0, ArrayList.class);
+            assert array1.equals(Arrays.asList(1L, 2L, 3L));
+
+            ArrayList<?> array2 = row.getAs(1, ArrayList.class);
+            assert array2.equals(Arrays.asList("a", "b"));
+
+            ArrayList<?> array3 = row.getAs(2, ArrayList.class);
+            assert array3.equals(Collections.singletonList(new Timestamp(31000000000L)));
+
+            ArrayList<?> array4 = row.getAs(3, ArrayList.class);
+            assert array4.equals(Collections.singletonList("[\n  1,\n  2\n]"));
+          } finally {
+            TimeZone.setDefault(oldTimeZone);
+          }
         },
         getSession());
   }

diff --git a/src/test/scala/com/snowflake/snowpark_test/RowSuite.scala b/src/test/scala/com/snowflake/snowpark_test/RowSuite.scala
@@ -372,30 +372,35 @@ class RowSuite extends SNTestBase {
 
   test("getAs with structured array") {
     structuredTypeTest {
-      TimeZone.setDefault(TimeZone.getTimeZone("US/Pacific"))
-
-      val query =
-        """SELECT
-          |    [1,2,3]::ARRAY(NUMBER) AS arr1,
-          |    ['a','b']::ARRAY(VARCHAR) AS arr2,
-          |    [parse_json(31000000)::timestamp_ntz]::ARRAY(TIMESTAMP_NTZ) AS arr3,
-          |    [[1,2]]::ARRAY(ARRAY) AS arr4
-          |""".stripMargin
-
-      val df = session.sql(query)
-      val row = df.collect()(0)
-
-      val array1 = row.getAs[Array[Object]](0)
-      assert(array1 sameElements Array(1, 2, 3))
-
-      val array2 = row.getAs[Array[Object]](1)
-      assert(array2 sameElements Array("a", "b"))
-
-      val array3 = row.getAs[Array[Object]](2)
-      assert(array3 sameElements Array(new Timestamp(31000000000L)))
-
-      val array4 = row.getAs[Array[Object]](3)
-      assert(array4 sameElements Array("[\n  1,\n  2\n]"))
+      val oldTimeZone = TimeZone.getDefault
+      try {
+        TimeZone.setDefault(TimeZone.getTimeZone("US/Pacific"))
+
+        val query =
+          """SELECT
+            |    [1,2,3]::ARRAY(NUMBER) AS arr1,
+            |    ['a','b']::ARRAY(VARCHAR) AS arr2,
+            |    [parse_json(31000000)::timestamp_ntz]::ARRAY(TIMESTAMP_NTZ) AS arr3,
+            |    [[1,2]]::ARRAY(ARRAY) AS arr4
+            |""".stripMargin
+
+        val df = session.sql(query)
+        val row = df.collect()(0)
+
+        val array1 = row.getAs[Array[Object]](0)
+        assert(array1 sameElements Array(1, 2, 3))
+
+        val array2 = row.getAs[Array[Object]](1)
+        assert(array2 sameElements Array("a", "b"))
+
+        val array3 = row.getAs[Array[Object]](2)
+        assert(array3 sameElements Array(new Timestamp(31000000000L)))
+
+        val array4 = row.getAs[Array[Object]](3)
+        assert(array4 sameElements Array("[\n  1,\n  2\n]"))
+      } finally {
+        TimeZone.setDefault(oldTimeZone)
+      }
     }
   }