diff --git a/tests/src/test/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContext.scala b/tests/src/test/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContext.scala
index 4bb3d0ca2ca..b0aec74feb7 100644
--- a/tests/src/test/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContext.scala
+++ b/tests/src/test/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContext.scala
@@ -29,7 +29,7 @@ import org.apache.spark.scheduler.TaskLocality
 import org.apache.spark.shuffle.FetchFailedException
 import org.apache.spark.util.{AccumulatorV2, TaskCompletionListener, TaskFailureListener}
 
-class MockTaskContext(taskAttemptId: Long, partitionId: Int) extends MockTaskContextBase {
+class MockTaskContext(taskAttemptId: Long, partitionId: Int) extends TaskContext {
 
   val listeners = new ListBuffer[TaskCompletionListener]
 
@@ -96,4 +96,10 @@ class MockTaskContext(taskAttemptId: Long, partitionId: Int) extends MockTaskCon
   def markTaskComplete(): Unit = {
     listeners.foreach(_.onTaskCompletion(this))
   }
+
+  /**
+   * This method was introduced in Spark-3.5.1. It's not shimmed and added to the common class by
+   * removing the override keyword.
+   */
+  def isFailed(): Boolean = false
 }
diff --git a/tests/src/test/spark311/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala b/tests/src/test/spark311/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala
deleted file mode 100644
index 17927026b9f..00000000000
--- a/tests/src/test/spark311/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala
+++ /dev/null
@@ -1,47 +0,0 @@
-/*
- * Copyright (c) 2023, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/*** spark-rapids-shim-json-lines
-{"spark": "311"}
-{"spark": "312"}
-{"spark": "313"}
-{"spark": "320"}
-{"spark": "321"}
-{"spark": "321cdh"}
-{"spark": "321db"}
-{"spark": "322"}
-{"spark": "323"}
-{"spark": "324"}
-{"spark": "330"}
-{"spark": "330cdh"}
-{"spark": "330db"}
-{"spark": "331"}
-{"spark": "332"}
-{"spark": "332cdh"}
-{"spark": "332db"}
-{"spark": "333"}
-{"spark": "340"}
-{"spark": "341"}
-{"spark": "341db"}
-{"spark": "342"}
-{"spark": "350"}
-spark-rapids-shim-json-lines ***/
-package org.apache.spark.sql.rapids.metrics.source
-
-import org.apache.spark.TaskContext
-
-abstract class MockTaskContextBase extends TaskContext {
-}
diff --git a/tests/src/test/spark351/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala b/tests/src/test/spark351/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala
deleted file mode 100644
index d39ad6a1add..00000000000
--- a/tests/src/test/spark351/scala/org/apache/spark/sql/rapids/metrics/source/MockTaskContextBase.scala
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Copyright (c) 2023, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/*** spark-rapids-shim-json-lines
-{"spark": "351"}
-spark-rapids-shim-json-lines ***/
-package org.apache.spark.sql.rapids.metrics.source
-
-import org.apache.spark.TaskContext
-
-abstract class MockTaskContextBase extends TaskContext {
-  override def isFailed(): Boolean = false
-}