From 07f37b207a66c3a7f420dc25113a06e519f16e6b Mon Sep 17 00:00:00 2001
From: DmitriyMX <dimon550@gmail.com>
Date: Wed, 16 Dec 2015 13:42:09 +0300
Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?=
 =?UTF-8?q?=D0=BD=20Formatter?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main/java/ru/dmitriymx/shell/Shell.java |  8 +++++-
 src/test/java/test/TestShell.java           | 31 +++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/main/java/ru/dmitriymx/shell/Shell.java b/src/main/java/ru/dmitriymx/shell/Shell.java
index de971e1..fa647de 100644
--- a/src/main/java/ru/dmitriymx/shell/Shell.java
+++ b/src/main/java/ru/dmitriymx/shell/Shell.java
@@ -21,6 +21,7 @@ public class Shell {
     protected ConsoleReader console;
     private CommandLoop commandLoop;
     private CommandCompleter commandCompleter;
+    private Formatter formatter;
     protected boolean run = false;
 
     public void start() throws IOException, InterruptedException {
@@ -70,7 +71,11 @@ public class Shell {
     }
 
     public void setFormatter(Formatter formatter) {
-        newErr.setFormatter(formatter);
+        if (newErr != null) {
+            newErr.setFormatter(formatter);
+        } else {
+            this.formatter = formatter;
+        }
     }
 
     /**
@@ -79,6 +84,7 @@ public class Shell {
     public void overrideSysOutErr() {
         sysErr = System.err;
         newErr = new ShellPrintStream(sysErr);
+        if (formatter != null) newErr.setFormatter(formatter);
         System.setErr(newErr);
         System.setOut(newErr);
     }
diff --git a/src/test/java/test/TestShell.java b/src/test/java/test/TestShell.java
index d4f09d8..5a3f533 100644
--- a/src/test/java/test/TestShell.java
+++ b/src/test/java/test/TestShell.java
@@ -1,12 +1,15 @@
 package test;
 
+import org.fusesource.jansi.Ansi;
 import org.junit.Before;
 import org.junit.Test;
+import ru.dmitriymx.shell.Formatter;
 import ru.dmitriymx.shell.Shell;
 
 import java.io.*;
 import java.util.logging.Logger;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 /**
@@ -85,4 +88,32 @@ public class TestShell {
         assertTrue(resultOut.contains("123"));
         assertTrue(resultOut.contains("Unknown command"));
     }
+
+    @Test
+    public void testFormatter() throws IOException, InterruptedException {
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        PrintStream shellout = new PrintStream(baos);
+        System.setOut(shellout);
+        System.setErr(shellout);
+
+        ByteArrayInputStream bais = new ByteArrayInputStream("qwe\nexit\n".getBytes());
+        System.setIn(bais);
+
+        Formatter formatter = s -> "FORMAT: [" + s + "]";
+
+        Shell shell = new Shell();
+        shell.setFormatter(formatter);
+        shell.overrideSysOutErr();
+        shell.start();
+
+        while (shell.isRunning()) {
+            Thread.sleep(1);
+        }
+
+        String resultOut = baos.toString();
+        returnSys();
+
+        assertTrue(resultOut.contains("FORMAT: ["));
+        assertTrue(resultOut.contains("Unknown command"));
+    }
 }
\ No newline at end of file