@Test
public void testOnDeviceAllocated() throws Exception {

  NvidiaGPUPluginForRuntimeV2 plugin = new NvidiaGPUPluginForRuntimeV2();
  Set<Device> allocatedDevices = new TreeSet<>();
  DeviceRuntimeSpec spec = plugin.onDevicesAllocated(allocatedDevices,
      YarnRuntimeType.RUNTIME_DEFAULT);
  Assert.assertNull(spec);
  // allocate one device
  allocatedDevices.add(Device.Builder.newInstance()
      .setId(0).setHealthy(true)
      .setBusID("00000000:04:00.0")
      .setDevPath("/dev/nvidia0")
      .setMajorNumber(195)
      .setMinorNumber(0).build());
  spec = plugin.onDevicesAllocated(allocatedDevices,
      YarnRuntimeType.RUNTIME_DOCKER);
  Assert.assertEquals("nvidia", spec.getContainerRuntime());
  Assert.assertEquals("0", spec.getEnvs().get("NVIDIA_VISIBLE_DEVICES"));
  // two device allowed
  allocatedDevices.add(Device.Builder.newInstance()
      .setId(0).setHealthy(true)
      .setBusID("00000000:82:00.0")
      .setDevPath("/dev/nvidia1")
      .setMajorNumber(195)
      .setMinorNumber(1).build());
  spec = plugin.onDevicesAllocated(allocatedDevices,
      YarnRuntimeType.RUNTIME_DOCKER);
  Assert.assertEquals("nvidia", spec.getContainerRuntime());
  Assert.assertEquals("0,1", spec.getEnvs().get("NVIDIA_VISIBLE_DEVICES"));
}