@Test
  public void shouldUpdateNumberOfActiveRemoteAgentsAfterAssigned() {
    AgentConfig agentConfig = AgentMother.remoteAgent();
    configHelper.addAgent(agentConfig);
    fixture.createPipelineWithFirstStageScheduled();
    buildAssignmentService.onTimer();

    int before = agentService.numberOfActiveRemoteAgents();
    Work work = buildAssignmentService.assignWorkToAgent(agent(agentConfig));
    assertThat(work, instanceOf(BuildWork.class));
    assertThat(agentService.numberOfActiveRemoteAgents(), is(before + 1));
  }
  @Test
  public void shouldCallForReregisterIfAgentInstanceIsNotRegistered() {
    AgentConfig agentConfig = AgentMother.remoteAgent();
    fixture.createPipelineWithFirstStageScheduled();
    AgentRuntimeInfo info =
        AgentRuntimeInfo.fromServer(agentConfig, true, "location", 1000000l, "OS");
    agentService.requestRegistration(info);

    assertThat(agentService.findAgent(info.getUUId()).isRegistered(), is(false));

    info.setCookie("cookie");
    agentRemoteHandler.process(agent, new Message(Action.ping, info));
    buildAssignmentService.onTimer();

    assertThat(agent.messages.size(), is(1));
    assertThat(agent.messages.get(0).getAction(), is(Action.reregister));
  }
  @Test
  public void shouldAssignMatchedJobToAgentsRegisteredInAgentRemoteHandler() {
    AgentConfig agentConfig = AgentMother.remoteAgent();
    configHelper.addAgent(agentConfig);
    fixture.createPipelineWithFirstStageScheduled();
    AgentRuntimeInfo info =
        AgentRuntimeInfo.fromServer(agentConfig, true, "location", 1000000l, "OS");
    info.setCookie("cookie");

    agentRemoteHandler.process(agent, new Message(Action.ping, info));

    int before = agentService.numberOfActiveRemoteAgents();

    buildAssignmentService.onTimer();

    assertThat(agent.messages.size(), is(1));
    assertThat(agent.messages.get(0).getData(), instanceOf(BuildWork.class));
    assertThat(agentService.numberOfActiveRemoteAgents(), is(before + 1));
  }
 @After
 public void teardown() throws Exception {
   goCache.clear();
   agentService.clearAll();
   fixture.onTearDown();
   dbHelper.onTearDown();
   configHelper.onTearDown();
   FileUtil.deleteFolder(goConfigService.artifactsDir());
   agentAssignment.clear();
   agentRemoteHandler.connectedAgents().clear();
 }
  @Test
  public void shouldNotAssignWorkToCanceledAgentsRegisteredInAgentRemoteHandler() {
    AgentConfig agentConfig = AgentMother.remoteAgent();
    configHelper.addAgent(agentConfig);
    fixture.createPipelineWithFirstStageScheduled();
    AgentRuntimeInfo info =
        AgentRuntimeInfo.fromServer(agentConfig, true, "location", 1000000l, "OS");
    info.setCookie("cookie");

    agentRemoteHandler.process(agent, new Message(Action.ping, info));

    AgentInstance agentInstance = agentService.findAgentAndRefreshStatus(info.getUUId());
    agentInstance.cancel();

    buildAssignmentService.onTimer();

    assertThat(
        "Should not assign work when agent status is Canceled", agent.messages.size(), is(0));
  }
 private AgentIdentifier agent(AgentConfig agentConfig) {
   agentService.sync(new Agents(agentConfig));
   agentService.approve(agentConfig.getUuid());
   return agentService.findAgent(agentConfig.getUuid()).getAgentIdentifier();
 }