Testing output streams in Java

Here's a handy little helper test class for unit testing output streams in Java:

public class AssertOutputStream extends ByteArrayOutputStream {
    private final byte[] expected;

    public AssertOutputStream(final byte[] expected) {
        if (null == expected) throw new NullPointerException();

        this.expected = expected;
    }

    @Override
    public void close()
            throws IOException {
        super.close();

        Assert.assertTrue(Arrays.equals(expected, toByteArray()));
    }
}

Just drop it into a unit test as the destination output stream along with an expected byte array to test that the expected bytes were written out. If you prefer a failure message, it is easy enough to add a member String to hold one for the call to Assert.assertTrue.

UPDATE: Following the suggestion of ThoughtWorker Paul Holser, here is an inproved version:

public class AssertOutputStream extends ByteArrayOutputStream {
    private final String message;
    private final byte[] expected;

    public AssertOutputStream(final byte[] expected) {
        this(null, expected);
    }

    public AssertOutputStream(final String message,
            final byte[] expected) {
        super(expected.length);

        this.message = message;
        this.expected = expected;
    }

    @Override
    public void close()
            throws IOException {
        super.close();

        Assert.assertTrue(message,
                Arrays.equals(expected, toByteArray()));
    }
}