Fix shuffle test helper construction

Author: gerashegalov

tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleClientSuite.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,7 +33,7 @@ class RapidsShuffleClientSuite extends RapidsShuffleTestHelper {
   def prepareBufferReceiveState(
       tableMeta: TableMeta,
       bounceBuffer: BounceBuffer): BufferReceiveState = {
-    val ptr = PendingTransferRequest(client, tableMeta, mockHandler)
+    val ptr = new PendingTransferRequest(client, tableMeta, mockHandler)
     spy(new BufferReceiveState(123L, bounceBuffer, Seq(ptr), () => {}))
   }
 
@@ -42,7 +42,7 @@ class RapidsShuffleClientSuite extends RapidsShuffleTestHelper {
       bounceBuffer: BounceBuffer): BufferReceiveState = {
 
     val ptrs = tableMetas.map { tm =>
-      PendingTransferRequest(client, tm, mockHandler)
+      new PendingTransferRequest(client, tm, mockHandler)
     }
 
     spy(new BufferReceiveState(123L, bounceBuffer, ptrs, () => {}))

tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleIteratorSuite.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -172,7 +172,7 @@ class RapidsShuffleIteratorSuite extends RapidsShuffleTestHelper {
       val ac = ArgumentCaptor.forClass(classOf[RapidsShuffleFetchHandler])
       when(mockTransport.makeClient(any())).thenReturn(client)
       doNothing().when(client).doFetch(any(), ac.capture())
-      val mockBuffer = RapidsShuffleHandle(mock[SpillableDeviceBufferHandle], null)
+      val mockBuffer = new RapidsShuffleHandle(mock[SpillableDeviceBufferHandle], null)
       when(mockBuffer.spillable.sizeInBytes).thenReturn(123L)
 
       val cb = new ColumnarBatch(Array.empty, 10)

tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServerSuite.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -57,7 +57,7 @@ class RapidsShuffleServerSuite extends RapidsShuffleTestHelper {
         fillBuffer(hostBuff)
         deviceBuffer.copyFromHostBuffer(hostBuff)
         val mockMeta = RapidsShuffleTestHelper.mockTableMeta(100000)
-        RapidsShuffleHandle(SpillableDeviceBufferHandle(deviceBuffer), mockMeta)
+        new RapidsShuffleHandle(SpillableDeviceBufferHandle(deviceBuffer), mockMeta)
       }
     }
     new MockRapidsShuffleRequestHandler(mockBuffers)
@@ -208,7 +208,7 @@ class RapidsShuffleServerSuite extends RapidsShuffleTestHelper {
       withResource(new RefCountedDirectByteBuffer(bb)) { _ =>
         val tableMeta = MetaUtils.buildTableMeta(1, 456, bb, 100)
         val testHandle = SpillableDeviceBufferHandle(DeviceMemoryBuffer.allocate(456))
-        val rapidsBuffer = RapidsShuffleHandle(testHandle, tableMeta)
+        val rapidsBuffer = new RapidsShuffleHandle(testHandle, tableMeta)
         when(mockRequestHandler.getShuffleHandle(ArgumentMatchers.eq(1)))
           .thenReturn(rapidsBuffer)
 
@@ -277,8 +277,8 @@ class RapidsShuffleServerSuite extends RapidsShuffleTestHelper {
       val ex = new IllegalStateException("something happened")
       when(mockHandleThatThrows.materialize()).thenThrow(ex)
 
-      val rapidsBuffer = RapidsShuffleHandle(mockHandle, tableMeta)
-      val rapidsBufferThatThrows = RapidsShuffleHandle(mockHandleThatThrows, tableMeta)
+      val rapidsBuffer = new RapidsShuffleHandle(mockHandle, tableMeta)
+      val rapidsBufferThatThrows = new RapidsShuffleHandle(mockHandleThatThrows, tableMeta)
 
       when(mockRequestHandler.getShuffleHandle(ArgumentMatchers.eq(1)))
         .thenReturn(rapidsBuffer)
@@ -359,7 +359,7 @@ class RapidsShuffleServerSuite extends RapidsShuffleTestHelper {
         val tableMeta = MetaUtils.buildTableMeta(tableId, 456, bb, 100)
         val rapidsBuffer = if (error) {
           val mockHandle = mock[SpillableDeviceBufferHandle]
-          val rapidsBuffer = RapidsShuffleHandle(mockHandle, tableMeta)
+          val rapidsBuffer = new RapidsShuffleHandle(mockHandle, tableMeta)
           when(mockHandle.sizeInBytes).thenReturn(tableMeta.bufferMeta().size())
           // mock an error with the copy
           when(rapidsBuffer.spillable.materialize())
@@ -369,7 +369,7 @@ class RapidsShuffleServerSuite extends RapidsShuffleTestHelper {
           rapidsBuffer
         } else {
           val testHandle = spy(SpillableDeviceBufferHandle(spy(DeviceMemoryBuffer.allocate(456))))
-          RapidsShuffleHandle(testHandle, tableMeta)
+          new RapidsShuffleHandle(testHandle, tableMeta)
         }
         when(mockRequestHandler.getShuffleHandle(ArgumentMatchers.eq(tableId)))
           .thenAnswer(_ => rapidsBuffer)

tests/src/test/spark330/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleTestHelper.scala

@@ -117,7 +117,7 @@ abstract class RapidsShuffleTestHelper
 
   def getSendBounceBuffer(size: Long): SendBounceBuffers = {
     val db = DeviceMemoryBuffer.allocate(size)
-    SendBounceBuffers(new BounceBuffer(db) {
+    new SendBounceBuffers(new BounceBuffer(db) {
       override def free(bb: BounceBuffer): Unit = {
         db.close()
       }

tests/src/test/spark340/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleTestHelper.scala

@@ -131,7 +131,7 @@ abstract class RapidsShuffleTestHelper
 
   def getSendBounceBuffer(size: Long): SendBounceBuffers = {
     val db = DeviceMemoryBuffer.allocate(size)
-    SendBounceBuffers(new BounceBuffer(db) {
+    new SendBounceBuffers(new BounceBuffer(db) {
       override def free(bb: BounceBuffer): Unit = {
         db.close()
       }