import java.io.*;
import java.util.Random;
import java.util.jar.JarOutputStream;
import java.util.zip.DeflaterOutputStream;
import java.util.zip.GZIPOutputStream;
import java.util.zip.InflaterOutputStream;
import java.util.zip.ZipOutputStream;
import java.util.zip.ZipEntry;

import org.testng.annotations.BeforeTest;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertThrows;


public class CloseInflaterDeflaterTest {

    // Number of bytes to write/read from Deflater/Inflater
    private static final int INPUT_LENGTH= 512;
    // OutputStream that will throw an exception during a write operation
    private static OutputStream outStream = new OutputStream() {
        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            throw new IOException();
        }
        @Override
        public void write(byte[] b) throws IOException {}
        @Override
        public void write(int b) throws IOException {}
    };

    // Input bytes for read/write operation
    private static byte[] inputBytes = new byte[INPUT_LENGTH];
    // Random function to add bytes to inputBytes
    private static Random rand = new Random();

    /**
     * DataProvider to specify whether to use close() or finish() of OutputStream
     *
     * @return Entry object indicating which method to use for closing OutputStream
     */
    @DataProvider
    public Object[][] testOutputStreams() {
        return new Object[][] {
                { true },
                { false },
        };
    }

    /**
     * DataProvider to specify on which outputstream closeEntry() has to be called
     *
     * @return Entry object returning either JarOutputStream or ZipOutputStream
     */
    @DataProvider
    public Object[][] testZipAndJar() throws IOException{
        return new Object[][] {
                { new JarOutputStream(outStream)},
                { new ZipOutputStream(outStream)},
        };
    }

    /**
     * Add inputBytes array with random bytes to write into OutputStream
     */
    @BeforeTest
    public void before_test()
    {
        rand.nextBytes(inputBytes);
    }

    /**
     * Test for infinite loop by writing bytes to closed GZIPOutputStream
     *
     * @param useCloseMethod indicates whether to use Close() or finish() method
     * @throws IOException if an error occurs
     */
    @Test(dataProvider = "testOutputStreams")
    public void testGZip(boolean useCloseMethod) throws IOException {
        GZIPOutputStream gzip = new GZIPOutputStream(outStream);
        gzip.write(inputBytes, 0, INPUT_LENGTH);
        assertThrows(IOException.class, () -> {
            // Close GZIPOutputStream
            if (useCloseMethod) {
                gzip.close();
            } else {
                gzip.finish();
            }
        });
        // Write on a closed GZIPOutputStream, IOException expected
        assertThrows(IOException.class , () -> gzip.write(inputBytes, 0, INPUT_LENGTH));

        assertThrows(IOException.class, () -> gzip.write(inputBytes, 0, 0));
    }

    /**
     * Test for infinite loop by writing bytes to closed DeflaterOutputStream
     *
     * @param useCloseMethod indicates whether to use Close() or finish() method
     * @throws IOException if an error occurs
     */
    @Test(dataProvider = "testOutputStreams")
    public void testDeflaterOutputStream(boolean useCloseMethod) throws IOException {
        DeflaterOutputStream def = new DeflaterOutputStream(outStream);
        assertThrows(IOException.class , () -> def.write(inputBytes, 0, INPUT_LENGTH));
        assertThrows(IOException.class, () -> {
            // Close DeflaterOutputStream
            if (useCloseMethod) {
                def.close();
            } else {
                def.finish();
            }
        });
        // Write on a closed DeflaterOutputStream, IOException expected
        assertThrows(IOException.class , () -> def.write(inputBytes, 0, INPUT_LENGTH));

        assertThrows(IOException.class, () -> def.write(inputBytes, 0, 0));
    }


    /**
     * Test for infinite loop by writing bytes to closed InflaterOutputStream
     *
     * Note: Disabling this test as it is failing intermittently.
     * @param useCloseMethod indicates whether to use Close() or finish() method
     * @throws IOException if an error occurs
     */
    @Test(dataProvider = "testOutputStreams",enabled=false)
    public void testInflaterOutputStream(boolean useCloseMethod) throws IOException {
        InflaterOutputStream inf = new InflaterOutputStream(outStream);
        assertThrows(IOException.class , () -> inf.write(inputBytes, 0, INPUT_LENGTH));
        assertThrows(IOException.class , () -> {
            // Close InflaterOutputStream
            if (useCloseMethod) {
                inf.close();
            } else {
                inf.finish();
            }
        });
        // Write on a closed InflaterOutputStream , IOException expected
        assertThrows(IOException.class , () -> inf.write(inputBytes, 0, INPUT_LENGTH));
    }

    /**
     * Test for infinite loop by writing bytes to closed ZipOutputStream/JarOutputStream
     *
     * @param zip will be the instance of either JarOutputStream or ZipOutputStream
     * @throws IOException if an error occurs
     */
    @Test(dataProvider = "testZipAndJar")
    public void testZipCloseEntry(ZipOutputStream zip) throws IOException {
        assertThrows(IOException.class , () -> zip.putNextEntry(new ZipEntry("")));
        zip.write(inputBytes, 0, INPUT_LENGTH);
        assertThrows(IOException.class , () -> zip.closeEntry());
        // Write on a closed ZipOutputStream, IOException expected
        assertThrows(IOException.class , () -> zip.write(inputBytes, 0, INPUT_LENGTH));
    }

}