Skip to content

Commit 94b621f

Browse files
sfc-gh-heshahsfc-gh-heshah
committed
SNOW-1062684: Finalize test changes after server-side support
1 parent af645d9 commit 94b621f

File tree

15 files changed

+114
-28
lines changed

15 files changed

+114
-28
lines changed

‎scripts/utils.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,8 @@ run_test_suites() {
7777
'++ 2.12.20 OtherTests:testOnly * -- -l SampleDataTest' \
7878
'++ 2.13.16 OtherTests:testOnly * -- -l SampleDataTest' \
7979
'++ 2.12.20 UDFTests:testOnly * -- -l SampleDataTest' \
80-
'++ 2.13.16 UDFTests:testOnly * -- -l SampleDataTest -l com.snowflake.snowpark.UDFPackageTest' \
80+
'++ 2.13.16 UDFTests:testOnly * -- -l SampleDataTest' \
8181
'++ 2.12.20 UDTFTests:testOnly * -- -l SampleDataTest' \
82-
'++ 2.13.16 UDTFTests:testOnly * -- -l SampleDataTest -l com.snowflake.snowpark.UDFPackageTest' \
82+
'++ 2.13.16 UDTFTests:testOnly * -- -l SampleDataTest' \
8383
+SprocTests:test
8484
}

‎src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import com.snowflake.snowpark.{
2525
}
2626

2727
import java.io._
28+
import com.snowflake.snowpark.internal.Utils.{ScalaCompatVersion, SnowparkPackageName}
2829
import com.snowflake.snowpark.types.{Geography, Geometry, Variant}
2930
import com.snowflake.snowpark_java.types.InternalUtils
3031
import com.snowflake.snowpark_java.udtf._
@@ -33,6 +34,10 @@ import scala.collection.{JavaConverters, mutable}
3334
import scala.collection.JavaConverters._
3435

3536
object JavaUtils {
37+
def snowparkScalaCompatVersion(): String = ScalaCompatVersion
38+
39+
def snowparkPackageName(): String = SnowparkPackageName
40+
3641
def session_setJavaAPI(builder: SessionBuilder): SessionBuilder =
3742
builder.setJavaAPI()
3843

‎src/test/java/com/snowflake/snowpark_test/JavaStoredProcedureSuite.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.snowflake.snowpark_test;
22

33
import com.snowflake.snowpark.SnowparkClientException;
4+
import com.snowflake.snowpark.internal.JavaUtils;
45
import com.snowflake.snowpark_java.Row;
56
import com.snowflake.snowpark_java.Session;
67
import com.snowflake.snowpark_java.StoredProcedure;
@@ -37,7 +38,9 @@ public void call() {
3738
+ "returns STRING\n"
3839
+ "language scala\n"
3940
+ "runtime_version=2.12\n"
40-
+ "packages=('com.snowflake:snowpark:latest')\n"
41+
+ "packages=('"
42+
+ JavaUtils.snowparkPackageName()
43+
+ ":latest')\n"
4144
+ "handler='Test.run'\n"
4245
+ "as\n"
4346
+ "$$\n"

‎src/test/java/com/snowflake/snowpark_test/UDFTestBase.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.snowflake.snowpark_test;
22

33
import com.snowflake.snowpark.TestUtils;
4+
import com.snowflake.snowpark.internal.JavaUtils;
45
import com.snowflake.snowpark_java.Session;
56

67
public abstract class UDFTestBase extends TestFunctions {
@@ -18,7 +19,11 @@ public Session getSession() {
1819
}
1920

2021
protected Session createSession() {
21-
return Session.builder().configFile(defaultProfile).create();
22+
Session newSession = Session.builder().configFile(defaultProfile).create();
23+
if (JavaUtils.snowparkScalaCompatVersion().equals("2.13")) {
24+
newSession.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect();
25+
}
26+
return newSession;
2227
}
2328

2429
protected void runQuery(String sql) {

‎src/test/scala/com/snowflake/snowpark/TestUtils.scala‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ object TestUtils extends Logging {
125125
sess.removeDependency(snClassDir)
126126
if (usePackages) {
127127
sess.removePackage(Utils.clientPackageName)
128-
sess.addPackage("com.snowflake:snowpark:latest")
128+
sess.addPackage(s"${Utils.SnowparkPackageName}:latest")
129129
} else {
130130
// Replace scoverage classes in classpath because they are instrumented.
131131
sess.addDependency(snClassDir.replace("scoverage-", ""))
@@ -147,10 +147,7 @@ object TestUtils extends Logging {
147147
classOf[org.scalatest.compatible.Assertion],
148148
classOf[org.scalactic.TripleEquals], // scalactic jar
149149
classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter],
150-
classOf[io.opentelemetry.sdk.trace.export.SpanExporter],
151-
classOf[scala.Product],
152-
classOf[scala.reflect.api.Universe],
153-
classOf[scala.tools.nsc.Global])
150+
classOf[io.opentelemetry.sdk.trace.export.SpanExporter])
154151
.flatMap(UDFClassPath.getPathForClass(_))
155152
.foreach(path => {
156153
val file = new File(path)
@@ -179,10 +176,7 @@ object TestUtils extends Logging {
179176
classOf[BeforeAndAfterAll], // scala test jar
180177
classOf[org.scalactic.TripleEquals], // scalactic jar
181178
classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter],
182-
classOf[io.opentelemetry.sdk.trace.export.SpanExporter],
183-
classOf[scala.Product],
184-
classOf[scala.reflect.api.Universe],
185-
classOf[scala.tools.nsc.Global])
179+
classOf[io.opentelemetry.sdk.trace.export.SpanExporter])
186180
.flatMap(UDFClassPath.getPathForClass(_))
187181
.foreach(path => {
188182
val file = new File(path)

‎src/test/scala/com/snowflake/snowpark/UDFClasspathSuite.scala‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package com.snowflake.snowpark
33
import java.io.File
44
import java.net.URLClassLoader
55
import com.snowflake.snowpark.internal.ScalaFunctions._toUdf
6-
import com.snowflake.snowpark.internal.Utils.clientPackageName
76
import com.snowflake.snowpark.internal.{UDFClassPath, UDXRegistrationHandler, Utils}
87
import org.mockito.ArgumentMatchers.any
98
import org.mockito.Mockito.{atLeastOnce, never, reset, spy, times, verify}
@@ -15,9 +14,19 @@ import scala.util.Random
1514
@UDFTest
1615
class UDFClasspathSuite extends SNTestBase {
1716

18-
override def beforeAll(): Unit = {
19-
super.beforeAll()
17+
override def beforeAll: Unit = {
18+
super.beforeAll
2019
TestUtils.addDepsToClassPath(session)
20+
if (Utils.ScalaCompatVersion == "2.13") {
21+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true")
22+
}
23+
}
24+
25+
override def afterAll: Unit = {
26+
if (Utils.ScalaCompatVersion == "2.13") {
27+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false")
28+
}
29+
super.afterAll
2130
}
2231

2332
test("Test that jars uploaded to different stages") {
@@ -97,8 +106,8 @@ class UDFClasspathSuite extends SNTestBase {
97106
val fixedPath = path.replace("scoverage-", "")
98107
// Remove snowpark jar from classpath, The code will catch error and add the path and retry
99108
newSession.removeDependency(fixedPath)
100-
newSession.removePackage("com.snowflake:snowpark:latest")
101-
newSession.removePackage(clientPackageName)
109+
newSession.removePackage(s"${Utils.SnowparkPackageName}:latest")
110+
newSession.removePackage(Utils.clientPackageName)
102111
val func = "func_" + Random.nextInt().abs
103112

104113
ignoreClassNotFoundForScoverageClasses {

‎src/test/scala/com/snowflake/snowpark/UDFInternalSuite.scala‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,16 @@ class UDFInternalSuite extends TestData {
2121
TestUtils.addDepsToClassPath(session, Some(stageName))
2222
TestUtils.addDepsToClassPath(newSession, Some(stageName))
2323
}
24+
if (Utils.ScalaCompatVersion == "2.13") {
25+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true")
26+
}
2427
}
2528

2629
override def afterAll: Unit = {
2730
dropStage(stageName)
31+
if (Utils.ScalaCompatVersion == "2.13") {
32+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false")
33+
}
2834
super.afterAll
2935
}
3036

@@ -46,7 +52,7 @@ class UDFInternalSuite extends TestData {
4652
verify(mockSession, times(1)).addDependency(path)
4753
}
4854
verify(mockSession, times(1)).removeDependency(path)
49-
verify(mockSession, times(1)).addPackage("com.snowflake:snowpark:latest")
55+
verify(mockSession, times(1)).addPackage(s"${Utils.SnowparkPackageName}:latest")
5056
}
5157

5258
test("Test permanent udf not failing back to upload jar", JavaStoredProcExclude) {
@@ -80,7 +86,7 @@ class UDFInternalSuite extends TestData {
8086
verify(mockSession, times(1)).addDependency(path)
8187
}
8288
verify(mockSession, times(1)).removeDependency(path)
83-
verify(mockSession, times(1)).addPackage("com.snowflake:snowpark:latest")
89+
verify(mockSession, times(1)).addPackage(s"${Utils.SnowparkPackageName}:latest")
8490
}
8591

8692
test("Test add version logic", JavaStoredProcExclude) {
@@ -189,7 +195,7 @@ class PackageUDFSuite extends UDFSuite {
189195
val snClassDir = UDFClassPath.getPathForClass(classOf[Session]).get
190196
session.removeDependency(snClassDir.replace("scoverage-", ""))
191197
session.removePackage(Utils.clientPackageName)
192-
session.addPackage("com.snowflake:snowpark:latest")
198+
session.addPackage(s"${Utils.SnowparkPackageName}:latest")
193199
}
194200

195201
override def afterAll: Unit = {
@@ -207,7 +213,7 @@ class PackageUDTFSuite extends UDTFSuite {
207213
val snClassDir = UDFClassPath.getPathForClass(classOf[Session]).get
208214
session.removeDependency(snClassDir.replace("scoverage-", ""))
209215
session.removePackage(Utils.clientPackageName)
210-
session.addPackage("com.snowflake:snowpark:latest")
216+
session.addPackage(s"${Utils.SnowparkPackageName}:latest")
211217
}
212218

213219
override def afterAll: Unit = {

‎src/test/scala/com/snowflake/snowpark/UDFRegistrationSuite.scala‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import com.snowflake.snowpark.internal.Utils.clientPackageName
44

55
import java.io.{BufferedOutputStream, File, FileOutputStream}
66
import java.nio.file.{Files, NoSuchFileException}
7-
import com.snowflake.snowpark.internal.{JavaUtils, UDFClassPath}
7+
import com.snowflake.snowpark.internal.{JavaUtils, UDFClassPath, Utils}
88

99
import scala.reflect.internal.util.BatchSourceFile
1010
import scala.reflect.io.{AbstractFile, VirtualDirectory}
@@ -17,12 +17,22 @@ class UDFRegistrationSuite extends SNTestBase with FileUtils {
1717
private val tempStage = this.getClass.getSimpleName + Random.alphanumeric.take(5).mkString("")
1818
private val stagePrefix = "udfJar"
1919

20-
override def beforeAll(): Unit = {
20+
override def beforeAll: Unit = {
2121
super.beforeAll()
2222
session.runQuery(s"create or replace temporary stage $tempStage")
2323
if (!isStoredProc(session)) {
2424
TestUtils.addDepsToClassPath(session)
2525
}
26+
if (Utils.ScalaCompatVersion == "2.13") {
27+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true")
28+
}
29+
}
30+
31+
override def afterAll: Unit = {
32+
if (Utils.ScalaCompatVersion == "2.13") {
33+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false")
34+
}
35+
super.afterAll
2636
}
2737

2838
test("Test that jar files are uploaded to stage correctly") {

‎src/test/scala/com/snowflake/snowpark/UDTFInternalSuite.scala‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package com.snowflake.snowpark
22

3-
import com.snowflake.snowpark.internal.ScalaFunctions
3+
import com.snowflake.snowpark.internal.{ScalaFunctions, Utils}
44
import com.snowflake.snowpark.types._
55
import com.snowflake.snowpark.udtf._
66

@@ -13,6 +13,16 @@ class UDTFInternalSuite extends SNTestBase {
1313
if (!isStoredProc(session)) {
1414
TestUtils.addDepsToClassPath(session)
1515
}
16+
if (Utils.ScalaCompatVersion == "2.13") {
17+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true")
18+
}
19+
}
20+
21+
override def afterAll: Unit = {
22+
if (Utils.ScalaCompatVersion == "2.13") {
23+
session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false")
24+
}
25+
super.afterAll
1626
}
1727

1828
test("Unit test for UDTF0") {

‎src/test/scala/com/snowflake/snowpark_test/PermanentUDFSuite.scala‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.snowflake.snowpark_test
22

33
import com.snowflake.snowpark.TestUtils.{removeFile, writeFile}
4+
import com.snowflake.snowpark.internal.Utils.ScalaCompatVersion
45
import com.snowflake.snowpark.functions._
56
import com.snowflake.snowpark._
67
import net.snowflake.client.jdbc.SnowflakeSQLException
@@ -30,12 +31,18 @@ class PermanentUDFSuite extends TestData {
3031
// In stored procs mode, there is only one session
3132
TestUtils.addDepsToClassPath(newSession, Some(stageName))
3233
}
34+
if (ScalaCompatVersion == "2.13") {
35+
session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect()
36+
}
3337
}
3438

3539
override def afterAll: Unit = {
3640
dropStage(stageName)
3741
removeFile(tempDirectory1.getAbsolutePath, session)
3842
removeFile(tempDirectory2.getAbsolutePath, session)
43+
if (ScalaCompatVersion == "2.13") {
44+
session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect()
45+
}
3946
super.afterAll
4047
}
4148

0 commit comments

Comments
 (0)