Add logic and tests for ListJobs

Author: gkiko10

databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java

@@ -3,6 +3,7 @@
 import com.databricks.sdk.core.ApiClient;
 import com.databricks.sdk.service.jobs.*;
 import java.util.Collection;
+import java.util.Iterator;
 
 public class JobsExt extends JobsAPI {
 
@@ -14,6 +15,34 @@ public JobsExt(JobsService mock) {
     super(mock);
   }
 
+  public Iterable<BaseJob> list(ListJobsRequest request) {
+    Iterable<BaseJob> jobsList =  super.list(request);
+
+    if (!request.getExpandTasks()) {
+      return jobsList;
+    }
+
+    Iterator<BaseJob> iterator = jobsList.iterator();
+    return () -> new Iterator<BaseJob>() {
+      @Override
+      public boolean hasNext() {
+        return iterator.hasNext();
+      }
+
+      @Override
+      public BaseJob next() {
+        BaseJob job = iterator.next();
+        GetJobRequest getJobRequest = new GetJobRequest().setJobId(job.getJobId());
+        Job fullJob = get(getJobRequest);
+        job.getSettings().setTasks(fullJob.getSettings().getTasks());
+        job.getSettings().setJobClusters(fullJob.getSettings().getJobClusters());
+        job.getSettings().setParameters(fullJob.getSettings().getParameters());
+        job.getSettings().setEnvironments(fullJob.getSettings().getEnvironments());
+        return job;
+      }
+    };
+  }
+
   /**
    * Wrap the {@code JobsApi.getRun} operation to retrieve paginated content without breaking the
    * response contract.

databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java

@@ -202,4 +202,109 @@ private void addJobEnvironments(Job job, String... environmentKeys) {
     }
     job.getSettings().setEnvironments(environments);
   }
+
+  @Test
+  public void testListJobsWithoutExpandTasks() {
+    JobsService service = Mockito.mock(JobsService.class);
+    BaseJob job1 = new BaseJob().setJobId(100L).setSettings(new JobSettings().setName("job1"));
+    BaseJob job2 = new BaseJob().setJobId(200L).setSettings(new JobSettings().setName("job2"));
+    BaseJob job3 = new BaseJob().setJobId(300L).setSettings(new JobSettings().setName("job3"));
+    BaseJob job4 = new BaseJob().setJobId(400L).setSettings(new JobSettings().setName("job4"));
+
+    List<BaseJob> jobsOnFirstPage = new ArrayList<>();
+    jobsOnFirstPage.add(job1);
+    jobsOnFirstPage.add(job2);
+    List<BaseJob> jobsOnSecondPage = new ArrayList<>();
+    jobsOnSecondPage.add(job3);
+    jobsOnSecondPage.add(job4);
+    when(service.list(any()))
+            .thenReturn(new ListJobsResponse().setJobs(jobsOnFirstPage))
+            .thenReturn(new ListJobsResponse().setJobs(jobsOnSecondPage));
+    JobsExt jobsExt = new JobsExt(service);
+
+    ListJobsRequest request = new ListJobsRequest().setExpandTasks(false);
+    Iterable<BaseJob> jobsList = jobsExt.list(request);
+
+    List<BaseJob> expectedJobsList = new ArrayList<>();
+    expectedJobsList.add(job1);
+    expectedJobsList.add(job2);
+    expectedJobsList.add(job3);
+    expectedJobsList.add(job4);
+    for (BaseJob job : jobsList) {
+      BaseJob expectedJob = expectedJobsList.stream()
+              .filter(e -> e.getJobId().equals(job.getJobId()))
+              .findFirst()
+              .orElse(null);
+      assertEquals(expectedJob, job);
+    }
+    verify(service, times(0)).get(any());
+  }
+
+  @Test
+  public void testListJobs() {
+    JobsService service = Mockito.mock(JobsService.class);
+    BaseJob job1 = new BaseJob().setJobId(100L).setSettings(new JobSettings().setName("job1"));
+    addTasks(job1, "job1_taskKey1", "job1_taskKey2");
+    BaseJob job2 = new BaseJob().setJobId(200L).setSettings(new JobSettings().setName("job2"));
+    addTasks(job2, "job2_taskKey1", "job2_taskKey2");
+
+    Job getJob1_page1 = new Job()
+            .setJobId(100L)
+            .setNextPageToken("job1_page2token")
+            .setSettings(new JobSettings().setName("job1"));
+    addTasks(getJob1_page1, "job1_taskKey1", "job1_taskKey2");
+    Job getJob1_page2 = new Job()
+            .setJobId(100L)
+            .setSettings(new JobSettings().setName("job1"));
+    addTasks(getJob1_page2, "job1_taskKey3", "job1_taskKey4");
+
+    Job getJob2_page1 = new Job()
+            .setJobId(200L)
+            .setNextPageToken("job2_page2token")
+            .setSettings(new JobSettings().setName("job2"));
+    addTasks(getJob2_page1, "job2_taskKey1", "job2_taskKey2");
+    Job getJob2_page2 = new Job()
+            .setJobId(200L)
+            .setSettings(new JobSettings().setName("job2"));
+    addTasks(getJob2_page2, "job2_taskKey3", "job2_taskKey4");
+
+    doReturn(getJob1_page1).when(service).get(argThat(request -> request != null && request.getJobId() == 100L && request.getPageToken() == null));
+    doReturn(getJob1_page2).when(service).get(argThat(request -> request != null && request.getJobId() == 100L && "job1_page2token".equals(request.getPageToken())));
+    doReturn(getJob2_page1).when(service).get(argThat(request -> request != null && request.getJobId() == 200L && request.getPageToken() == null));
+    doReturn(getJob2_page2).when(service).get(argThat(request -> request != null && request.getJobId() == 200L && "job2_page2token".equals(request.getPageToken())));
+
+    List<BaseJob> jobsOnFirstPage = new ArrayList<>();
+    jobsOnFirstPage.add(job1);
+    jobsOnFirstPage.add(job2);
+    ListJobsResponse response = new ListJobsResponse().setJobs(jobsOnFirstPage);
+    when(service.list(any())).thenReturn(response);
+    JobsExt jobsExt = new JobsExt(service);
+
+    ListJobsRequest request = new ListJobsRequest().setExpandTasks(true);
+    Iterable<BaseJob> jobsList = jobsExt.list(request);
+
+    BaseJob expectedJob1 = new BaseJob().setJobId(100L).setSettings(new JobSettings().setName("job1"));
+    addTasks(expectedJob1, "job1_taskKey1", "job1_taskKey2", "job1_taskKey3", "job1_taskKey4");
+    BaseJob expectedJob2 = new BaseJob().setJobId(200L).setSettings(new JobSettings().setName("job2"));
+    addTasks(expectedJob2, "job2_taskKey1", "job2_taskKey2", "job2_taskKey3", "job2_taskKey4");
+    List<BaseJob> expectedJobsList = new ArrayList<>();
+    expectedJobsList.add(expectedJob1);
+    expectedJobsList.add(expectedJob2);
+    for (BaseJob job : jobsList) {
+      BaseJob expectedJob = expectedJobsList.stream()
+              .filter(e -> e.getJobId().equals(job.getJobId()))
+              .findFirst()
+              .orElse(null);
+      assertEquals(expectedJob, job);
+    }
+    verify(service, times(4)).get(any());
+  }
+
+  private void addTasks(BaseJob job, String... taskKeys) {
+    Collection<Task> tasks = new ArrayList<>();
+    for (String taskKey : taskKeys) {
+      tasks.add(new Task().setTaskKey(taskKey));
+    }
+    job.getSettings().setTasks(tasks);
+  }
 }