@Test
public void testPrimaryWorkerComponentNumWorkersIsTwo() throws IOException {

  testCommons.yarnConfig.set("hadoop.registry.dns.domain-name", "testDomain");
  testCommons = new ComponentTestCommons(TensorFlowRole.PRIMARY_WORKER);
  testCommons.setupTensorFlow();
  TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters();
  parameters.setWorkerResource(testCommons.resource);
  parameters.setName("testJobName");
  parameters.setNumWorkers(2);
  parameters.setWorkerDockerImage("testWorkerDockerImage");
  TensorFlowWorkerComponent workerComponent =
      createWorkerComponent(parameters);
  Component component = workerComponent.createComponent();
  assertEquals(1L, (long) component.getNumberOfContainers());
  // If the dependencies are upgraded to hadoop 3.3.0.
  // yarn.service.container-state-report-as-service-state can be replaced
  // with CONTAINER_STATE_REPORT_AS_SERVICE_STATE
  verifyCommons(component, ImmutableMap.of(
      "yarn.service.container-state-report-as-service-state", "true"));
}