Do output encapsulation in a try/with

Move the shell output encapsulation so that we don't double-wrap on a
premature exit.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=219136896

java/google/registry/tools/ShellCommand.java

@@ -145,6 +145,80 @@ public class ShellCommand implements Command {
     return this;
   }
 
+  private static class OutputEncapsulator {
+    private PrintStream orgStdout;
+    private PrintStream orgStderr;
+
+    private EncapsulatingOutputStream encapsulatedOutputStream = null;
+    private EncapsulatingOutputStream encapsulatedErrorStream = null;
+
+    private Exception error;
+
+    private OutputEncapsulator() {
+      orgStdout = System.out;
+      orgStderr = System.err;
+      encapsulatedOutputStream = new EncapsulatingOutputStream(System.out, "out: ");
+      encapsulatedErrorStream = new EncapsulatingOutputStream(System.out, "err: ");
+      System.setOut(new PrintStream(encapsulatedOutputStream));
+      System.setErr(new PrintStream(encapsulatedErrorStream));
+    }
+
+    void setError(Exception e) {
+      error = e;
+    }
+
+    private void restoreOriginalStreams() {
+      try {
+        encapsulatedOutputStream.dumpLastLine();
+        encapsulatedErrorStream.dumpLastLine();
+        System.setOut(orgStdout);
+        System.setErr(orgStderr);
+        if (error != null) {
+          emitFailure(error);
+        } else {
+          emitSuccess();
+        }
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    /**
+     * Emit a success command separator.
+     *
+     * <p>Dumps the last line of output prior to doing this.
+     */
+    private void emitSuccess() {
+      System.out.println(SUCCESS);
+      System.out.flush();
+    }
+
+    /**
+     * Emit a failure message obtained from the throwable.
+     *
+     * <p>Dumps the last line of output prior to doing this.
+     */
+    private void emitFailure(Throwable e) {
+      System.out.println(
+          FAILURE
+              + e.getClass().getName()
+              + " "
+              + e.getMessage().replace("\\", "\\\\").replace("\n", "\\n"));
+    }
+
+    /** Run "func" with output encapsulation. */
+    static void run(CommandRunner runner, String[] args) {
+      OutputEncapsulator encapsulator = new OutputEncapsulator();
+      try {
+        runner.run(args);
+      } catch (Exception e) {
+        encapsulator.setError(e);
+      } finally {
+        encapsulator.restoreOriginalStreams();
+      }
+    }
+  }
+
   /** Run the shell until the user presses "Ctrl-D". */
   @Override
   public void run() {
@@ -154,23 +228,6 @@ public class ShellCommand implements Command {
     String line;
     DateTime lastTime = clock.nowUtc();
     while ((line = getLine()) != null) {
-      PrintStream orgStdout = null;
-      PrintStream orgStderr = null;
-      EncapsulatingOutputStream encapsulatedOutputStream = null;
-      EncapsulatingOutputStream encapsulatedErrorStream = null;
-
-
-      // Wrap standard output and error if requested. We have to do so here in run because the flags
-      // haven't been processed in the constructor.
-      if (encapsulateOutput) {
-        orgStdout = System.out;
-        orgStderr = System.err;
-        encapsulatedOutputStream = new EncapsulatingOutputStream(System.out, "out: ");
-        encapsulatedErrorStream = new EncapsulatingOutputStream(System.out, "err: ");
-        System.setOut(new PrintStream(encapsulatedOutputStream));
-        System.setErr(new PrintStream(encapsulatedErrorStream));
-      }
-
       // Make sure we're not idle for too long. Only relevant when we're "extra careful"
       if (!dontExitOnIdle
           && beExtraCareful
@@ -184,27 +241,17 @@ public class ShellCommand implements Command {
       if (lineArgs.length == 0) {
         continue;
       }
-      Exception lastError = null;
+
+      // Wrap standard output and error if requested. We have to do so here in run because the flags
+      // haven't been processed in the constructor.
+      if (encapsulateOutput) {
+        OutputEncapsulator.run(runner, lineArgs);
+      } else {
         try {
           runner.run(lineArgs);
         } catch (Exception e) {
-        lastError = e;
           System.err.println("Got an exception:\n" + e);
         }
-      try {
-        if (encapsulatedOutputStream != null) {
-          encapsulatedOutputStream.dumpLastLine();
-          encapsulatedErrorStream.dumpLastLine();
-          System.setOut(orgStdout);
-          System.setErr(orgStderr);
-          if (lastError == null) {
-            emitSuccess();
-          } else {
-            emitFailure(lastError);
-          }
-        }
-      } catch (IOException e) {
-        throw new RuntimeException(e);
       }
     }
     if (!encapsulateOutput) {
@@ -243,29 +290,6 @@ public class ShellCommand implements Command {
     return resultBuilder.build().toArray(new String[0]);
   }
 
-  /**
-   * Emit a success command separator.
-   *
-   * <p>Dumps the last line of output prior to doing this.
-   */
-  private void emitSuccess() {
-    System.out.println(SUCCESS);
-    System.out.flush();
-  }
-
-  /**
-   * Emit a failure message obtained from the throwable.
-   *
-   * <p>Dumps the last line of output prior to doing this.
-   */
-  private void emitFailure(Throwable e) {
-    System.out.println(
-        FAILURE
-            + e.getClass().getName()
-            + " "
-            + e.getMessage().replace("\\", "\\\\").replace("\n", "\\n"));
-  }
-
   @VisibleForTesting
   static class JCommanderCompletor implements Completor {
 

javatests/google/registry/tools/ShellCommandTest.java

@@ -53,6 +53,9 @@ public class ShellCommandTest {
   PrintStream orgStdout;
   PrintStream orgStderr;
 
+  ByteArrayOutputStream stdout;
+  ByteArrayOutputStream stderr;
+
   public ShellCommandTest() {}
 
   @Before
@@ -269,14 +272,7 @@ public class ShellCommandTest {
   @Test
   public void testEncapsulatedOutput_command() throws Exception {
     RegistryToolEnvironment.ALPHA.setup();
-
-    // capture output (have to do this before the shell command is created)
-    ByteArrayOutputStream stdout = new ByteArrayOutputStream();
-    ByteArrayOutputStream stderr = new ByteArrayOutputStream();
-    System.setOut(new PrintStream(stdout));
-    System.setErr(new PrintStream(stderr));
-    System.setIn(new ByteArrayInputStream("command1\n".getBytes(UTF_8)));
-
+    captureOutput();
     ShellCommand shellCommand =
         new ShellCommand(
             args -> {
@@ -296,6 +292,33 @@ public class ShellCommandTest {
                 + "SUCCESS\n");
   }
 
+  @Test
+  public void testEncapsulatedOutput_noCommand() throws Exception {
+    captureOutput();
+    ShellCommand shellCommand =
+        createShellCommand(
+            args -> {
+              System.out.println("first line");
+            },
+            Duration.ZERO,
+            "",
+            "do something");
+    shellCommand.encapsulateOutput = true;
+    shellCommand.run();
+    assertThat(stderr.toString()).isEmpty();
+    assertThat(stdout.toString())
+        .isEqualTo("out: first line\nSUCCESS\n");
+  }
+
+  void captureOutput() {
+    // capture output (have to do this before the shell command is created)
+    stdout = new ByteArrayOutputStream();
+    stderr = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(stdout));
+    System.setErr(new PrintStream(stderr));
+    System.setIn(new ByteArrayInputStream("command1\n".getBytes(UTF_8)));
+  }
+
   @Parameters(commandDescription = "Test command")
   static class TestCommand implements Command {
     enum OrgType {