@Test
public void testPytorchServiceSpec() throws IOException {

  testCommons = new ComponentTestCommons(TensorFlowRole.PRIMARY_WORKER);
  testCommons.setupTensorFlow();
  PyTorchRunJobParameters parameters = new PyTorchRunJobParameters();
  parameters.setWorkerResource(testCommons.resource);
  parameters.setName("testJobName");
  parameters.setNumWorkers(1);
  parameters.setWorkerLaunchCmd("testWorkerLaunchCommand");
  MockClientContext mockClientContext = new MockClientContext(DEFAULT_JOB_NAME);
  FileSystemOperations fsOperations =
      new FileSystemOperations(mockClientContext);
  HadoopEnvironmentSetup hadoopEnv =
      new HadoopEnvironmentSetup(mockClientContext, fsOperations);
  PyTorchLaunchCommandFactory launchCommandFactory =
      new PyTorchLaunchCommandFactory(hadoopEnv, parameters,
          new Configuration());
  Localizer localizer = new Localizer(fsOperations,
      mockClientContext.getRemoteDirectoryManager(), parameters);
  PyTorchServiceSpec pyTorchServiceSpec = new PyTorchServiceSpec(parameters,
      mockClientContext, fsOperations, launchCommandFactory, localizer);
  ServiceWrapper serviceWrapper = pyTorchServiceSpec.create();
  Service service = serviceWrapper.getService();
  assertNotNull("Service must not be null!", service);
  List<Component> components = service.getComponents();
  assertEquals("Number of components is not correct!", 1, components.size());
  Component component = components.get(0);
  assertEquals(1L, (long) component.getNumberOfContainers());
  assertEquals("./run-PRIMARY_WORKER.sh", component.getLaunchCommand());
}