Examples with ElasticAgentPluginRegistry com.thoughtworks.go.plugin.access.elastic.ElasticAgentPluginRegistry used on opensource projects

Example 1 with ElasticAgentPluginRegistry

use of com.thoughtworks.go.plugin.access.elastic.ElasticAgentPluginRegistry in project gocd by gocd.

the class ScheduleServiceIntegrationTest method shouldSendCreateAgentRequestToPluginForTheJobRequiringElasticAgentWhichHasBeenRescheduled.

@Test
public // This test is written to verify fix for github issue https://github.com/gocd/gocd/issues/6615
void shouldSendCreateAgentRequestToPluginForTheJobRequiringElasticAgentWhichHasBeenRescheduled() {
    String pluginId = "ecs";
    ClusterProfile clusterProfile = new ClusterProfile("cluster_profile", pluginId);
    ElasticProfile elasticAgentProfile = new ElasticProfile("elastic_agent_profile", "cluster_profile");
    String ephemeralAutoRegisterKey = "auto_registry_key";
    // Mock elastic agent extension
    ElasticAgentPluginRegistry elasticAgentPluginRegistry = mock(ElasticAgentPluginRegistry.class);
    when(elasticAgentPluginRegistry.shouldAssignWork(any(), any(), any(), any(), any(), any())).thenReturn(true);
    when(elasticAgentPluginRegistry.has(any())).thenReturn(true);
    when(ephemeralAutoRegisterKeyService.autoRegisterKey()).thenReturn(ephemeralAutoRegisterKey);
    elasticAgentPluginService.setElasticAgentPluginRegistry(elasticAgentPluginRegistry);
    elasticAgentPluginService.setEphemeralAutoRegisterKeyService(ephemeralAutoRegisterKeyService);
    // Mock CreateAgentQueueHandler to verify create agent call was sent to the plugin
    CreateAgentQueueHandler createAgentQueueHandler = mock(CreateAgentQueueHandler.class);
    elasticAgentPluginService.setCreateAgentQueue(createAgentQueueHandler);
    // add a cluster profile and elastic agent profile to the config.
    GoConfigDao goConfigDao = configHelper.getGoConfigDao();
    goConfigDao.loadForEditing().getElasticConfig().getClusterProfiles().add(clusterProfile);
    goConfigDao.loadForEditing().getElasticConfig().getProfiles().add(elasticAgentProfile);
    // create 2 elastic agents for ecs plugin
    Agent agent = AgentMother.elasticAgent();
    agent.setElasticAgentId("elastic-agent-id-1");
    agent.setElasticPluginId(pluginId);
    Agent agentConfig2 = AgentMother.elasticAgent();
    agentConfig2.setElasticAgentId("elastic-agent-id-2");
    agentConfig2.setElasticPluginId(pluginId);
    dbHelper.addAgent(agent);
    dbHelper.addAgent(agentConfig2);
    // define a job in the config requiring elastic agent
    PipelineConfig pipelineToBeAdded = PipelineConfigMother.createPipelineConfigWithStages(UUID.randomUUID().toString(), "s1");
    pipelineToBeAdded.first().getJobs().first().setElasticProfileId("elastic_agent_profile");
    PipelineConfig pipelineConfig = configHelper.addPipeline(pipelineToBeAdded);
    // trigger the pipeline
    Pipeline pipeline = dbHelper.schedulePipeline(pipelineConfig, forceBuild(pipelineConfig), "Bob", new TimeProvider(), Collections.singletonMap("elastic_agent_profile", elasticAgentProfile), Collections.singletonMap("cluster_profile", clusterProfile));
    buildAssignmentService.onTimer();
    // verify no agents are building
    AgentInstance agent1 = agentService.findAgent(agent.getUuid());
    assertFalse(agent1.isBuilding());
    AgentInstance agent2 = agentService.findAgent(agentConfig2.getUuid());
    assertFalse(agent2.isBuilding());
    // assign the current job to elastic agent 1
    Work work = buildAssignmentService.assignWorkToAgent(agent(agent));
    assertThat(work, instanceOf(BuildWork.class));
    assertTrue(agent1.isBuilding());
    // reschedule the job
    scheduleService.rescheduleJob(pipeline.getFirstStage().getFirstJob());
    // when the job is rescheduled, the current job instance is marked as rescheduled and a job with new build id is created for it.
    long existingRescheduledBuildId = pipeline.getFirstStage().getFirstJob().getId();
    // newly scheduled build will have the next counter.
    long newlyScheduledBuildId = pipeline.getFirstStage().getFirstJob().getId() + 1;
    // existing job instance, which is rescheduled
    JobInstance existingRescheduledInstance = jobInstanceService.buildById(existingRescheduledBuildId);
    // newly created job instance
    JobInstance newlyScheduledInstance = jobInstanceService.buildById(newlyScheduledBuildId);
    assertThat(existingRescheduledInstance.getState(), is(JobState.Rescheduled));
    assertThat(newlyScheduledInstance.getState(), is(JobState.Scheduled));
    // verify the newly created instance's job plan include elastic profile and cluster profile
    assertThat(jobInstanceDao.loadPlan(newlyScheduledBuildId).getClusterProfile(), is(clusterProfile));
    assertThat(jobInstanceDao.loadPlan(newlyScheduledBuildId).getElasticProfile(), is(elasticAgentProfile));
    JobPlan jobPlanOfRescheduledInstance = jobInstanceDao.loadPlan(existingRescheduledBuildId);
    // invoke jobStatusListener to simulate first agent reporting job completed scenario
    jobStatusListener.onMessage(new JobStatusMessage(jobPlanOfRescheduledInstance.getIdentifier(), JobState.Completed, jobPlanOfRescheduledInstance.getAgentUuid()));
    // verify the newly created instance's job plan include elastic profile and cluster profile even after the first agent reports job completion
    assertThat(jobInstanceDao.loadPlan(newlyScheduledBuildId).getClusterProfile(), is(clusterProfile));
    assertThat(jobInstanceDao.loadPlan(newlyScheduledBuildId).getElasticProfile(), is(elasticAgentProfile));
    elasticAgentPluginService.createAgentsFor(Collections.emptyList(), Collections.singletonList(jobInstanceDao.loadPlan(newlyScheduledBuildId)));
    // verify create agent request was sent to the plugin
    CreateAgentMessage message = new CreateAgentMessage(ephemeralAutoRegisterKey, null, elasticAgentProfile, clusterProfile, jobPlanOfRescheduledInstance.getIdentifier());
    verify(createAgentQueueHandler, times(1)).post(message, 110000L);
}