gistfile1.java

/**
 *
 */
package lesson.test.util;

import java.io.BufferedOutputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.io.PrintStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

/**
 * @author 
 *
 */
public class StdioTester {
    private InputStream _saved_in;
    private PrintStream _saved_out;


    /**
     * save in/out.
     */
    public void setUp() {
        _saved_in = System.in;
        _saved_out = System.out;
    }


    /**
     * restore in/out.
     */
    public void tearDown() {
        System.setIn(_saved_in);
        System.setOut(_saved_out);
    }

    /**
     * set stdin.
     *
     * @param buf
     */
    public void setIn(byte[] buf) {
        _saved_in = System.in;
        System.setIn(new ByteArrayInputStream(buf));
    }

    /**
     * wrap stdout to get stream printed out.
     *
     * @return ByteArrayOutputStream to read from stdout
     */
    public final ByteArrayOutputStream setOut() {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        BufferedOutputStream bos = new BufferedOutputStream(baos);
        System.setOut(new PrintStream(bos));
        return baos;
    }

    public final <T> String setInOut(String str, Class<T> cls) {
        setIn(str.getBytes());
        ByteArrayOutputStream out = setOut();

        try {
            Method method = cls.getMethod("main", new Class[]{ String[].class });
            method.invoke(cls, new Object[]{ new String[0] });
        } catch (NoSuchMethodException | SecurityException
                | IllegalAccessException | IllegalArgumentException
                | InvocationTargetException e) {
            e.printStackTrace();
        }
        System.out.flush();

        return out.toString();
    }
}