package com.mockobjects.io;

import com.mockobjects.ExpectationCounter;
import com.mockobjects.ExpectationValue;
import com.mockobjects.Verifiable;
import com.mockobjects.util.Verifier;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;

public class MockOutputStream extends OutputStream implements Verifiable {
    private ExpectationValue myWriteWasCalled = new ExpectationValue("MockOutputStream.writeWasCalled");
    private ExpectationCounter myCloseCalls = new ExpectationCounter("MockOutputStream.close()");
    private ExpectationCounter myFlushCalls = new ExpectationCounter("MockOutputStream.flush()");
    private ByteArrayOutputStream myBuffer = new ByteArrayOutputStream();
    private boolean shouldThrowException = false;

    public MockOutputStream() {
        super();
    }

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

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

    public void flush() throws IOException {
        myFlushCalls.inc();
    }

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

    /**
     * Added in response to Bug report 546661.
     */
    public byte[] getContentsAsByteArray() {
        return myBuffer.toByteArray();
    }

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

    public void setExpectedFlushCalls(int flushCall) {
        myFlushCalls.setExpected(flushCall);
    }

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

    public void setupThrowIOException(boolean throwException) {
        shouldThrowException = throwException;
    }

    public void verify() {
        Verifier.verifyObject(this);
    }

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