fix test

snowflakedb · Jul 23, 2024 · 4bbede1 · 4bbede1
1 parent 05d1baa
commit 4bbede1
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 14 deletions.
diff --git a/src/main/scala/com/snowflake/snowpark/UDTFRegistration.scala b/src/main/scala/com/snowflake/snowpark/UDTFRegistration.scala
@@ -229,7 +229,7 @@ class UDTFRegistration(session: Session) extends Logging {
       "UDTFRegistration",
       funcName,
       execName,
-      s"${UDXRegistrationHandler.className}.${UDXRegistrationHandler.methodName}",
+      UDXRegistrationHandler.udtfClassName,
       execFilePath,
       0)(func)
   }

diff --git a/src/main/scala/com/snowflake/snowpark/internal/OpenTelemetry.scala b/src/main/scala/com/snowflake/snowpark/internal/OpenTelemetry.scala
@@ -43,7 +43,7 @@ object OpenTelemetry extends Logging {
       className,
       funcName,
       execName,
-      s"${UDXRegistrationHandler.className}.${UDXRegistrationHandler.methodName}",
+      UDXRegistrationHandler.udtfClassName,
       execFilePath,
       stackOffset + 2)(func.get())
   }

diff --git a/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala b/src/main/scala/com/snowflake/snowpark/internal/UDXRegistrationHandler.scala
@@ -34,7 +34,7 @@ object UDXRegistrationHandler {
   val className = "SnowUDF"
   // Method name for generated Java code
   val methodName = "compute"
-  private val udtfClassName = className
+  val udtfClassName = "SnowparkGeneratedUDTF"
 }
 
 class UDXRegistrationHandler(session: Session) extends Logging {

diff --git a/src/test/java/com/snowflake/snowpark_test/JavaUDXOpenTelemetrySuite.java b/src/test/java/com/snowflake/snowpark_test/JavaUDXOpenTelemetrySuite.java
@@ -53,17 +53,17 @@ public void udf() {
   public void udtf() {
     String className = "snow.snowpark.UDTFRegistration";
     getSession().udtf().registerTemporary(new MyJavaUDTFOf0());
-    checkUdfSpan(className, "registerTemporary", "", "");
+    checkUdtfSpan(className, "registerTemporary", "", "");
     String funcName = randomFunctionName();
     String funcName2 = randomFunctionName();
     getSession().udtf().registerTemporary(funcName, new MyJavaUDTFOf0());
-    checkUdfSpan(className, "registerTemporary", funcName, "");
+    checkUdtfSpan(className, "registerTemporary", funcName, "");
 
     String stageName = randomStageName();
     try {
       createStage(stageName, false);
       getSession().udtf().registerPermanent(funcName2, new MyJavaUDTFOf2(), stageName);
-      checkUdfSpan(className, "registerPermanent", funcName2, stageName);
+      checkUdtfSpan(className, "registerPermanent", funcName2, stageName);
     } finally {
       getSession().sql("drop function " + funcName2 + "(int, int)").collect();
       dropStage(stageName);
@@ -104,4 +104,18 @@ private void checkUdfSpan(
         "SnowUDF.compute",
         execFilePath);
   }
+
+  private void checkUdtfSpan(
+      String className, String funcName, String execName, String execFilePath) {
+    StackTraceElement[] stack = Thread.currentThread().getStackTrace();
+    StackTraceElement file = stack[2];
+    checkSpan(
+        className,
+        funcName,
+        "JavaUDXOpenTelemetrySuite.java",
+        file.getLineNumber() - 1,
+        execName,
+        "SnowparkGeneratedUDTF",
+        execFilePath);
+  }
 }
diff --git a/src/test/scala/com/snowflake/snowpark/TestUtils.scala b/src/test/scala/com/snowflake/snowpark/TestUtils.scala
@@ -145,9 +145,7 @@ object TestUtils extends Logging {
       classOf[BeforeAndAfterAll], // scala test jar
       classOf[org.scalactic.TripleEquals], // scalactic jar
       classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter],
-      classOf[io.opentelemetry.sdk.trace.export.SpanExporter],
-      classOf[scala.collection.Iterator[_]],
-      classOf[scala.reflect.api.TypeTags])
+      classOf[io.opentelemetry.sdk.trace.export.SpanExporter])
       .flatMap(UDFClassPath.getPathForClass(_))
       .foreach(path => {
         val file = new File(path)
@@ -178,8 +176,7 @@ object TestUtils extends Logging {
       classOf[BeforeAndAfterAll], // scala test jar
       classOf[org.scalactic.TripleEquals], // scalactic jar
       classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter],
-      classOf[io.opentelemetry.sdk.trace.export.SpanExporter],
-      classOf[scala.Function0[_]])
+      classOf[io.opentelemetry.sdk.trace.export.SpanExporter])
       .flatMap(UDFClassPath.getPathForClass(_))
       .foreach(path => {
         val file = new File(path)

diff --git a/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala b/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala
@@ -44,17 +44,17 @@ class UdxOpenTelemetrySuite extends OpenTelemetryEnabled {
       override def outputSchema(): StructType = StructType(StructField("sum", IntegerType))
     }
     session.udtf.registerTemporary(new MyUDTF0())
-    checkUdfSpan(className, "registerTemporary", "", "")
+    checkUdtfSpan(className, "registerTemporary", "", "")
     val udtfName = randomFunctionName()
     session.udtf.registerTemporary(udtfName, new MyUDTF0)
-    checkUdfSpan(className, "registerTemporary", udtfName, "")
+    checkUdtfSpan(className, "registerTemporary", udtfName, "")
 
     val stageName: String = randomName()
     val udtfName2 = randomFunctionName()
     try {
       createStage(stageName, isTemporary = false)
       session.udtf.registerPermanent(udtfName2, new MyUDTF0(), stageName)
-      checkUdfSpan(className, "registerPermanent", udtfName2, stageName)
+      checkUdtfSpan(className, "registerPermanent", udtfName2, stageName)
     } finally {
       runQuery(s"drop function $udtfName2()", session)
       dropStage(stageName)
@@ -97,4 +97,21 @@ class UdxOpenTelemetrySuite extends OpenTelemetryEnabled {
       "SnowUDF.compute",
       execFilePath)
   }
+
+  def checkUdtfSpan(
+      className: String,
+      funcName: String,
+      execName: String,
+      execFilePath: String): Unit = {
+    val stack = Thread.currentThread().getStackTrace
+    val file = stack(2) // this file
+    checkSpan(
+      className,
+      funcName,
+      "UdxOpenTelemetrySuite.scala",
+      file.getLineNumber - 1,
+      execName,
+      "SnowparkGeneratedUDTF",
+      execFilePath)
+  }
 }