Modified patch from Rob Spoor.

Author: Gary Gregory

src/test/java/org/apache/commons/io/output/TeeOutputStreamTest.java

@@ -17,51 +17,36 @@
 package org.apache.commons.io.output;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 
-import org.junit.Assert;
+import org.apache.commons.io.testtools.YellOnCloseOutputStream;
 import org.junit.Test;
 
-/**
+/**On
+ * JUnit Test Case for {@link TeeOutputStream}.
  */
-
 public class TeeOutputStreamTest {
 
-    private static class ExceptionOnCloseByteArrayOutputStream extends ByteArrayOutputStream {
-
-        @Override
-        public void close() throws IOException {
-            throw new IOException();
-        }
-    }
-
-    private static class RecordCloseByteArrayOutputStream extends ByteArrayOutputStream {
-
-        boolean closed;
-
-        @Override
-        public void close() throws IOException {
-            super.close();
-            closed = true;
-        }
-    }
-
     /**
      * Tests that the branch {@code OutputStream} is closed when closing the main {@code OutputStream} throws an
      * exception on {@link TeeOutputStream#close()}.
      */
     @Test
-    public void testCloseBranchIOException() {
-        final ByteArrayOutputStream badOs = new ExceptionOnCloseByteArrayOutputStream();
-        final RecordCloseByteArrayOutputStream goodOs = new RecordCloseByteArrayOutputStream();
+    public void testIOExceptionOnCloseBranch() throws IOException {
+        final OutputStream badOs = new YellOnCloseOutputStream();
+        final ByteArrayOutputStream goodOs = mock(ByteArrayOutputStream.class);
         final TeeOutputStream tos = new TeeOutputStream(goodOs, badOs);
         try {
             tos.close();
-            Assert.fail("Expected " + IOException.class.getName());
+            fail("Expected " + IOException.class.getName());
         } catch (final IOException e) {
-            Assert.assertTrue(goodOs.closed);
+            verify(goodOs).close();
         }
     }
 
@@ -70,15 +55,15 @@ public void testCloseBranchIOException() {
      * exception on {@link TeeOutputStream#close()}.
      */
     @Test
-    public void testCloseMainIOException() {
-        final ByteArrayOutputStream badOs = new ExceptionOnCloseByteArrayOutputStream();
-        final RecordCloseByteArrayOutputStream goodOs = new RecordCloseByteArrayOutputStream();
+    public void testIOExceptionOnClose() throws IOException {
+        final OutputStream badOs = new YellOnCloseOutputStream();
+        final ByteArrayOutputStream goodOs = mock(ByteArrayOutputStream.class);
         final TeeOutputStream tos = new TeeOutputStream(badOs, goodOs);
         try {
             tos.close();
-            Assert.fail("Expected " + IOException.class.getName());
+            fail("Expected " + IOException.class.getName());
         } catch (final IOException e) {
-            Assert.assertTrue(goodOs.closed);
+            verify(goodOs).close();
         }
     }