Fix to failing invalid device test

Summary:
The InitializationFailsWithInvalidDeviceId test was failing because of
an error in the cuda mock.

Reviewed By: MittalMakwana

Differential Revision: D85512707

fbshipit-source-id: 48123587312f8d9838e948844b0f74ab26c2c56b

Author: pavanbalaji

comms/torchcomms/ncclx/tests/unit/cpp/TorchCommNCCLXTest.cpp

@@ -141,6 +141,23 @@ TEST_F(TorchCommNCCLXTest, InitializationFailsWithInvalidDeviceId) {
     comm->setCudaApi(cuda_mock_);
     comm->setNcclApi(nccl_mock_);
 
+    // Mock getDeviceCount to return a valid device count (needed for rank %
+    // device_count)
+    EXPECT_CALL(*cuda_mock_, getDeviceCount(_))
+        .WillOnce(DoAll(SetArgPointee<0>(1), Return(cudaSuccess)));
+
+    // Mock malloc for barrier buffer allocation in bootstrap constructor
+    EXPECT_CALL(*cuda_mock_, malloc(_, sizeof(float)))
+        .Times(2)
+        .WillRepeatedly(DoAll(
+            SetArgPointee<0>(reinterpret_cast<void*>(0x1000)),
+            Return(cudaSuccess)));
+
+    // Mock free for barrier buffer deallocation in bootstrap destructor
+    EXPECT_CALL(*cuda_mock_, free(reinterpret_cast<void*>(0x1000)))
+        .Times(2)
+        .WillRepeatedly(Return(cudaSuccess));
+
     // Mock CUDA API to be called with device ID 0, since the boostrap
     // logic will assign a device ID in this case based on the rank.
     EXPECT_CALL(*cuda_mock_, setDevice(0))

comms/torchcomms/ncclx/tests/unit/cpp/mocks/CudaMock.cpp

@@ -26,6 +26,9 @@ void CudaMock::setupDefaultBehaviors() {
           SetArgPointee<1>(std::numeric_limits<int>::max()),
           Return(cudaSuccess)));
 
+  ON_CALL(*this, getDeviceCount(_))
+      .WillByDefault(DoAll(SetArgPointee<0>(1), Return(cudaSuccess)));
+
   // Stream management - return success by default
   ON_CALL(*this, getStreamPriorityRange(_, _))
       .WillByDefault(DoAll(