001/*
002 * Copyright 2012-2017 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.boot.test.rule;
018
019import java.io.ByteArrayOutputStream;
020import java.io.IOException;
021import java.io.OutputStream;
022import java.io.PrintStream;
023import java.util.ArrayList;
024import java.util.List;
025
026import org.hamcrest.Matcher;
027import org.junit.Assert;
028import org.junit.rules.TestRule;
029import org.junit.runner.Description;
030import org.junit.runners.model.Statement;
031
032import org.springframework.boot.ansi.AnsiOutput;
033import org.springframework.boot.ansi.AnsiOutput.Enabled;
034
035import static org.hamcrest.Matchers.allOf;
036
037/**
038 * JUnit {@code @Rule} to capture output from System.out and System.err.
039 *
040 * @author Phillip Webb
041 * @author Andy Wilkinson
042 * @since 1.4.0
043 */
044public class OutputCapture implements TestRule {
045
046        private CaptureOutputStream captureOut;
047
048        private CaptureOutputStream captureErr;
049
050        private ByteArrayOutputStream copy;
051
052        private List<Matcher<? super String>> matchers = new ArrayList<>();
053
054        @Override
055        public Statement apply(Statement base, Description description) {
056                return new Statement() {
057                        @Override
058                        public void evaluate() throws Throwable {
059                                captureOutput();
060                                try {
061                                        base.evaluate();
062                                }
063                                finally {
064                                        try {
065                                                if (!OutputCapture.this.matchers.isEmpty()) {
066                                                        String output = OutputCapture.this.toString();
067                                                        Assert.assertThat(output, allOf(OutputCapture.this.matchers));
068                                                }
069                                        }
070                                        finally {
071                                                releaseOutput();
072                                        }
073                                }
074                        }
075                };
076        }
077
078        protected void captureOutput() {
079                AnsiOutputControl.get().disableAnsiOutput();
080                this.copy = new ByteArrayOutputStream();
081                this.captureOut = new CaptureOutputStream(System.out, this.copy);
082                this.captureErr = new CaptureOutputStream(System.err, this.copy);
083                System.setOut(new PrintStream(this.captureOut));
084                System.setErr(new PrintStream(this.captureErr));
085        }
086
087        protected void releaseOutput() {
088                AnsiOutputControl.get().enabledAnsiOutput();
089                System.setOut(this.captureOut.getOriginal());
090                System.setErr(this.captureErr.getOriginal());
091                this.copy = null;
092        }
093
094        /**
095         * Discard all currently accumulated output.
096         */
097        public void reset() {
098                this.copy.reset();
099        }
100
101        public void flush() {
102                try {
103                        this.captureOut.flush();
104                        this.captureErr.flush();
105                }
106                catch (IOException ex) {
107                        // ignore
108                }
109        }
110
111        @Override
112        public String toString() {
113                flush();
114                return this.copy.toString();
115        }
116
117        /**
118         * Verify that the output is matched by the supplied {@code matcher}. Verification is
119         * performed after the test method has executed.
120         * @param matcher the matcher
121         */
122        public void expect(Matcher<? super String> matcher) {
123                this.matchers.add(matcher);
124        }
125
126        private static class CaptureOutputStream extends OutputStream {
127
128                private final PrintStream original;
129
130                private final OutputStream copy;
131
132                CaptureOutputStream(PrintStream original, OutputStream copy) {
133                        this.original = original;
134                        this.copy = copy;
135                }
136
137                @Override
138                public void write(int b) throws IOException {
139                        this.copy.write(b);
140                        this.original.write(b);
141                        this.original.flush();
142                }
143
144                @Override
145                public void write(byte[] b) throws IOException {
146                        write(b, 0, b.length);
147                }
148
149                @Override
150                public void write(byte[] b, int off, int len) throws IOException {
151                        this.copy.write(b, off, len);
152                        this.original.write(b, off, len);
153                }
154
155                public PrintStream getOriginal() {
156                        return this.original;
157                }
158
159                @Override
160                public void flush() throws IOException {
161                        this.copy.flush();
162                        this.original.flush();
163                }
164
165        }
166
167        /**
168         * Allow AnsiOutput to not be on the test classpath.
169         */
170        private static class AnsiOutputControl {
171
172                public void disableAnsiOutput() {
173                }
174
175                public void enabledAnsiOutput() {
176                }
177
178                public static AnsiOutputControl get() {
179                        try {
180                                Class.forName("org.springframework.boot.ansi.AnsiOutput");
181                                return new AnsiPresentOutputControl();
182                        }
183                        catch (ClassNotFoundException ex) {
184                                return new AnsiOutputControl();
185                        }
186                }
187
188        }
189
190        private static class AnsiPresentOutputControl extends AnsiOutputControl {
191
192                @Override
193                public void disableAnsiOutput() {
194                        AnsiOutput.setEnabled(Enabled.NEVER);
195                }
196
197                @Override
198                public void enabledAnsiOutput() {
199                        AnsiOutput.setEnabled(Enabled.DETECT);
200                }
201
202        }
203
204}