package com.mockobjects.servlet;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import javax.servlet.ServletOutputStream;
import com.mockobjects.ExpectationCounter;
import com.mockobjects.ExpectationValue;

public class MockServletOutputStream extends ServletOutputStream {
    private ExpectationValue myWriteCalled = new ExpectationValue("MockServletOutputStream.write()");
    private boolean myThrowException = false;
    private ExpectationCounter myCloseCallCount = new ExpectationCounter("MockServletOutputstream.close()");
    private ByteArrayOutputStream myBuffer;

    public MockServletOutputStream() {
        super();
        setupClearContents();
    }

    public void setExpectedCloseCalls(int closeCall) {
        myCloseCallCount.setExpected(closeCall);
    }

    public void setExpectingWriteCalls(boolean expectingWriteCall) {
        myWriteCalled.setExpected(expectingWriteCall);
    }

    public void setThrowIOException(boolean throwException) {
        myThrowException = throwException;
    }

    public void close() throws IOException {
        myCloseCallCount.inc();
    }

    public String toString() {
        return getContents();
    }

    public void write(int b) throws IOException {
        myWriteCalled.setActual(true);
        if (myThrowException)
            throw new IOException("Test IOException generated by request");
        myBuffer.write(b);
    }

    public void setupClearContents () {
        myBuffer = new ByteArrayOutputStream();
    }

    public String getContents() {
        return myBuffer.toString();
    }

    public void verify() {
        myWriteCalled.verify();
        myCloseCallCount.verify();
    }
}
