diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala
index 10fcaf620dfe9..c40149368b055 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala
@@ -44,7 +44,6 @@ class InvokeTargetClass extends Serializable {
   def filterInt(e: Any): Any = e.asInstanceOf[Int] > 0
   def filterPrimitiveInt(e: Int): Boolean = e > 0
   def binOp(e1: Int, e2: Double): Double = e1 + e2
-  def mapFunc(e: Any): Tuple2[Any, Any] = (e, e)
 }
 
 class InvokeTargetSubClass extends InvokeTargetClass {
@@ -203,20 +202,6 @@ class ObjectExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       Invoke(funcSubObj, "binOp", DoubleType, inputSum), 0.75, InternalRow.apply(1, 0.25))
   }
 
-  test("SPARK-31854: Invoke in MapElementsExec should not propagate null") {
-    val targetObject = new InvokeTargetClass
-    val funcClass = classOf[InvokeTargetClass]
-    val funcObj = Literal.create(targetObject, ObjectType(funcClass))
-    val inputInt = Seq(BoundReference(0, ObjectType(classOf[Any]), true))
-    val outputType = ObjectType(classOf[(Any, Any)])
-    val inputRow = InternalRow.fromSeq(Seq(null.asInstanceOf[java.lang.Integer]))
-    val createExpr = (propagateNull: Boolean) => {
-      Invoke(funcObj, "mapFunc", outputType, inputInt, propagateNull)
-    }
-    checkObjectExprEvaluation(createExpr(true), null, inputRow)
-    checkObjectExprEvaluation(createExpr(false), (null, null), inputRow)
-  }
-
   test("SPARK-23593: InitializeJavaBean should support interpreted execution") {
     val list = new java.util.LinkedList[Int]()
     list.add(1)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index bd09f46fcf819..06600c1e4b1d7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -1918,7 +1918,6 @@ class DatasetSuite extends QueryTest
   }
 
   test("SPARK-31854: Invoke in MapElementsExec should not propagate null") {
-    spark.conf.set("spark.sql.codegen.wholeStage", false)
     Seq("true", "false").foreach { wholeStage =>
       withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage) {
         val ds = Seq(1.asInstanceOf[Integer], null.asInstanceOf[Integer]).toDS()