Home | History | Annotate | Download | only in testing 
      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 
     17 package vogar.testing;
     18 
     19 import com.google.common.base.Joiner;
     20 import java.io.ByteArrayOutputStream;
     21 import java.io.IOException;
     22 import java.io.PrintStream;
     23 import java.io.UnsupportedEncodingException;
     24 import java.util.EnumMap;
     25 import java.util.EnumSet;
     26 import org.junit.rules.TestRule;
     27 import org.junit.runner.Description;
     28 import 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  */
     34 public 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