1/*
2 * Copyright (C) 2015 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17package vogar.testing;
18
19import com.google.common.base.Joiner;
20import java.io.ByteArrayOutputStream;
21import java.io.IOException;
22import java.io.PrintStream;
23import java.io.UnsupportedEncodingException;
24import java.util.EnumMap;
25import java.util.EnumSet;
26import org.junit.rules.TestRule;
27import org.junit.runner.Description;
28import org.junit.runners.model.Statement;
29
30/**
31 * A {@link TestRule} that will intercept content written to {@link System#out} and/or
32 * {@link System#err} and collate it for use by the test.
33 */
34public class InterceptOutputStreams implements TestRule {
35
36    /**
37     * The streams that can be intercepted.
38     */
39    public enum Stream {
40        OUT {
41            @Override
42            PrintStream get() {
43                return System.out;
44            }
45
46            @Override
47            void set(PrintStream stream) {
48                System.setOut(stream);
49            }
50        },
51        ERR {
52            @Override
53            PrintStream get() {
54                return System.err;
55            }
56
57            @Override
58            void set(PrintStream stream) {
59                System.setErr(stream);
60            }
61        };
62
63        abstract PrintStream get();
64
65        abstract void set(PrintStream stream);
66    }
67
68    /**
69     * The streams to intercept.
70     */
71    private final EnumSet<Stream> streams;
72    private final EnumMap<Stream, State> streams2State;
73
74    /**
75     * The streams to intercept.
76     */
77    public InterceptOutputStreams(Stream... streams) {
78        this.streams = EnumSet.of(streams[0], streams);
79        streams2State = new EnumMap<>(Stream.class);
80    }
81
82    /**
83     * Get the intercepted contents for the stream.
84     * @param stream the stream whose contents are required.
85     * @return the intercepted contents.
86     * @throws IllegalStateException if the stream contents are not being intercepted (in which
87     *     case the developer needs to add {@code stream} to the constructor parameters), or if the
88     *     test is not actually running at the moment.
89     */
90    public String contents(Stream stream) {
91        if (!streams.contains(stream)) {
92            EnumSet<Stream> extra = streams.clone();
93            extra.add(stream);
94            String message = "Not intercepting " + stream + " output, try:\n"
95                    + "    new " + InterceptOutputStreams.class.getSimpleName() + "("
96                    + Joiner.on(", ").join(extra)
97                    + ")";
98            throw new IllegalStateException(message);
99        }
100
101        State state = streams2State.get(stream);
102        if (state == null) {
103            throw new IllegalStateException(
104                    "Attempting to access stream contents outside the test");
105        }
106
107        return state.contents();
108    }
109
110    @Override
111    public Statement apply(final Statement base, Description description) {
112        return new Statement() {
113            @Override
114            public void evaluate() throws Throwable {
115                for (Stream stream : streams) {
116                    State state = new State(stream);
117                    streams2State.put(stream, state);
118                }
119
120                try {
121                    base.evaluate();
122                } finally {
123                    for (State state : streams2State.values()) {
124                        state.reset();
125                    }
126                    streams2State.clear();
127                }
128            }
129        };
130    }
131
132    private static class State {
133        private final PrintStream original;
134        private final ByteArrayOutputStream baos;
135        private final Stream stream;
136
137        State(Stream stream) throws IOException {
138            this.stream = stream;
139            original = stream.get();
140            baos = new ByteArrayOutputStream();
141            stream.set(new PrintStream(baos, true, "UTF-8"));
142        }
143
144        String contents() {
145            try {
146                return baos.toString("UTF-8");
147            } catch (UnsupportedEncodingException e) {
148                throw new RuntimeException(e);
149            }
150        }
151
152        void reset() {
153            stream.set(original);
154        }
155    }
156}
157