diff --git a/src/test/scala/com/databricks/labs/mosaic/expressions/geometry/ST_ZBehaviors.scala b/src/test/scala/com/databricks/labs/mosaic/expressions/geometry/ST_ZBehaviors.scala
index 3fde07b09..bd160a895 100644
--- a/src/test/scala/com/databricks/labs/mosaic/expressions/geometry/ST_ZBehaviors.scala
+++ b/src/test/scala/com/databricks/labs/mosaic/expressions/geometry/ST_ZBehaviors.scala
@@ -23,6 +23,7 @@ trait ST_ZBehaviors extends MosaicSpatialQueryTest {
             .getWKTRowsDf()
             .orderBy("id")
             .select("wkt")
+            .where("wkt SUBSTRING(wkt, 1, 5) = 'POINT'")
             .as[String]
             .collect()
             .map(wkt => mc.getGeometryAPI.geometry(wkt, "WKT"))
@@ -46,7 +47,7 @@ trait ST_ZBehaviors extends MosaicSpatialQueryTest {
         sqlResult.zip(expected).foreach { case (l, r) => l.equals(r) shouldEqual true }
     }
 
-    def stxCodegen(mosaicContext: MosaicContext): Unit = {
+    def stzCodegen(mosaicContext: MosaicContext): Unit = {
         spark.sparkContext.setLogLevel("FATAL")
         val mc = mosaicContext
         val sc = spark