Merge pull request #194 from spotify/drewc/json-helios-status

[cli] status command with --host --json only return jobs on those hosts

Author: drewcsillag

helios-system-tests/src/main/java/com/spotify/helios/system/QueryFailureTest.java

@@ -84,6 +84,6 @@ public void testJobStatusHostFilter() throws Exception {
     assertEquals(CreateJobResponse.Status.OK, created.getStatus());
 
     final String result = cli("status", "--host", "framazama");
-    assertThat(result, containsString("matched no hosts"));
+    assertThat(result, containsString("matched no jobs"));
   }
 }

helios-tools/src/main/java/com/spotify/helios/cli/command/JobStatusCommand.java

@@ -73,6 +73,11 @@ public JobStatusCommand(final Subparser parser) {
         .help("Print full hostnames, job and container id's.");
   }
 
+  interface HostStatusDisplayer {
+    void matchedStatus(final JobStatus jobStatus, final Iterable<String> matchingHosts,
+                       final Map<String, TaskStatus> taskStatuses);
+  }
+
   @Override
   int run(Namespace options, HeliosClient client, PrintStream out, final boolean json)
       throws ExecutionException, InterruptedException {
@@ -103,12 +108,58 @@ int run(Namespace options, HeliosClient client, PrintStream out, final boolean j
     statuses.putAll(allAsMap(futures));
 
     if (json) {
-      out.println(Json.asPrettyStringUnchecked(statuses));
+      showJsonStatuses(out, hostPattern, jobIds, statuses);
       return 0;
     }
 
     final JobStatusTable table = jobStatusTable(out, full);
 
+    final boolean noHostMatchedEver = showStatusesForHosts(hostPattern, jobIds, statuses,
+      new HostStatusDisplayer() {
+        @Override
+        public void matchedStatus(JobStatus jobStatus, Iterable<String> matchingHosts,
+                                  Map<String, TaskStatus> taskStatuses) {
+          displayTask(full, table, jobStatus.getJob().getId(), jobStatus, taskStatuses,
+              matchingHosts);
+        }
+    });
+
+    if (noHostMatchedEver) {
+      out.printf("host pattern %s matched no jobs%n", hostPattern);
+      return 1;
+    }
+
+    table.print();
+
+    return 0;
+  }
+
+  private void showJsonStatuses(PrintStream out, final String hostPattern, final Set<JobId> jobIds,
+      final Map<JobId, JobStatus> statuses) {
+    if (Strings.isNullOrEmpty(hostPattern)) {
+      out.println(Json.asPrettyStringUnchecked(statuses));
+      return;
+    }
+
+    final Map<JobId, JobStatus> returnStatuses = Maps.newTreeMap();
+    showStatusesForHosts(hostPattern, jobIds, statuses, new HostStatusDisplayer() {
+      @Override
+      public void matchedStatus(JobStatus jobStatus, Iterable<String> matchingHosts,
+                                Map<String, TaskStatus> taskStatuses) {
+        for (final String host : matchingHosts) {
+          final Map<String, Deployment> deployments = jobStatus.getDeployments();
+          final Deployment deployment = (deployments == null) ? null : deployments.get(host);
+          if (deployment != null) {
+            returnStatuses.put(jobStatus.getJob().getId(), jobStatus);
+          }
+        }
+      }
+    });
+    out.println(Json.asPrettyStringUnchecked(returnStatuses));
+  }
+
+  private boolean showStatusesForHosts(final String hostPattern, final Set<JobId> jobIds,
+      final Map<JobId, JobStatus> statuses, final HostStatusDisplayer statusDisplayer) {
     boolean noHostMatchedEver = true;
 
     for (final JobId jobId : Ordering.natural().sortedCopy(jobIds)) {
@@ -133,21 +184,20 @@ int run(Namespace options, HeliosClient client, PrintStream out, final boolean j
         noHostMatchedEver = false;
       }
 
-      for (final String host : matchingHosts) {
-        final Map<String, Deployment> deployments = jobStatus.getDeployments();
-        final TaskStatus ts = taskStatuses.get(host);
-        final Deployment deployment = (deployments == null) ? null : deployments.get(host);
-        table.task(jobId, formatHostname(full, host), ts, deployment);
-      }
-    }
+      statusDisplayer.matchedStatus(jobStatus, matchingHosts, taskStatuses);
 
-    if (noHostMatchedEver) {
-      out.printf("host pattern %s matched no hosts%n", hostPattern);
-      return 1;
     }
+    return noHostMatchedEver;
+  }
 
-    table.print();
-
-    return 0;
+  private void displayTask(final boolean full, final JobStatusTable table, final JobId jobId,
+      final JobStatus jobStatus, final Map<String, TaskStatus> taskStatuses,
+      final Iterable<String> matchingHosts) {
+    for (final String host : matchingHosts) {
+      final Map<String, Deployment> deployments = jobStatus.getDeployments();
+      final TaskStatus ts = taskStatuses.get(host);
+      final Deployment deployment = (deployments == null) ? null : deployments.get(host);
+      table.task(jobId, formatHostname(full, host), ts, deployment);
+    }
   }
 }