1/*
2 * Copyright (C) 2014 Square, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16package okio;
17
18import java.io.IOException;
19import java.io.InputStream;
20import java.util.Arrays;
21import java.util.Random;
22import java.util.zip.Deflater;
23import java.util.zip.Inflater;
24import java.util.zip.InflaterInputStream;
25import org.junit.Test;
26
27import static org.junit.Assert.assertEquals;
28import static org.junit.Assert.fail;
29
30public final class DeflaterSinkTest {
31  @Test public void deflateWithClose() throws Exception {
32    OkBuffer data = new OkBuffer();
33    String original = "They're moving in herds. They do move in herds.";
34    data.writeUtf8(original);
35    OkBuffer sink = new OkBuffer();
36    DeflaterSink deflaterSink = new DeflaterSink(sink, new Deflater());
37    deflaterSink.write(data, data.size());
38    deflaterSink.close();
39    OkBuffer inflated = inflate(sink);
40    assertEquals(original, inflated.readUtf8(inflated.size()));
41  }
42
43  @Test public void deflateWithSyncFlush() throws Exception {
44    String original = "Yes, yes, yes. That's why we're taking extreme precautions.";
45    OkBuffer data = new OkBuffer();
46    data.writeUtf8(original);
47    OkBuffer sink = new OkBuffer();
48    DeflaterSink deflaterSink = new DeflaterSink(sink, new Deflater());
49    deflaterSink.write(data, data.size());
50    deflaterSink.flush();
51    OkBuffer inflated = inflate(sink);
52    assertEquals(original, inflated.readUtf8(inflated.size()));
53  }
54
55  @Test public void deflateWellCompressed() throws IOException {
56    String original = repeat('a', 1024 * 1024);
57    OkBuffer data = new OkBuffer();
58    data.writeUtf8(original);
59    OkBuffer sink = new OkBuffer();
60    DeflaterSink deflaterSink = new DeflaterSink(sink, new Deflater());
61    deflaterSink.write(data, data.size());
62    deflaterSink.close();
63    OkBuffer inflated = inflate(sink);
64    assertEquals(original, inflated.readUtf8(inflated.size()));
65  }
66
67  @Test public void deflatePoorlyCompressed() throws IOException {
68    ByteString original = randomBytes(1024 * 1024);
69    OkBuffer data = new OkBuffer();
70    data.write(original);
71    OkBuffer sink = new OkBuffer();
72    DeflaterSink deflaterSink = new DeflaterSink(sink, new Deflater());
73    deflaterSink.write(data, data.size());
74    deflaterSink.close();
75    OkBuffer inflated = inflate(sink);
76    assertEquals(original, inflated.readByteString(inflated.size()));
77  }
78
79  @Test public void multipleSegmentsWithoutCompression() throws IOException {
80    OkBuffer buffer = new OkBuffer();
81    Deflater deflater = new Deflater();
82    deflater.setLevel(Deflater.NO_COMPRESSION);
83    DeflaterSink deflaterSink = new DeflaterSink(buffer, deflater);
84    int byteCount = Segment.SIZE * 4;
85    deflaterSink.write(new OkBuffer().writeUtf8(repeat('a', byteCount)), byteCount);
86    deflaterSink.close();
87    assertEquals(repeat('a', byteCount), inflate(buffer).readUtf8(byteCount));
88  }
89
90  /**
91   * This test deflates a single segment of without compression because that's
92   * the easiest way to force close() to emit a large amount of data to the
93   * underlying sink.
94   */
95  @Test public void closeWithExceptionWhenWritingAndClosing() throws IOException {
96    MockSink mockSink = new MockSink();
97    mockSink.scheduleThrow(0, new IOException("first"));
98    mockSink.scheduleThrow(1, new IOException("second"));
99    Deflater deflater = new Deflater();
100    deflater.setLevel(Deflater.NO_COMPRESSION);
101    DeflaterSink deflaterSink = new DeflaterSink(mockSink, deflater);
102    deflaterSink.write(new OkBuffer().writeUtf8(repeat('a', Segment.SIZE)), Segment.SIZE);
103    try {
104      deflaterSink.close();
105      fail();
106    } catch (IOException expected) {
107      assertEquals("first", expected.getMessage());
108    }
109    mockSink.assertLogContains("close()");
110  }
111
112  /**
113   * Uses streaming decompression to inflate {@code deflated}. The input must
114   * either be finished or have a trailing sync flush.
115   */
116  private OkBuffer inflate(OkBuffer deflated) throws IOException {
117    InputStream deflatedIn = deflated.inputStream();
118    Inflater inflater = new Inflater();
119    InputStream inflatedIn = new InflaterInputStream(deflatedIn, inflater);
120    OkBuffer result = new OkBuffer();
121    byte[] buffer = new byte[8192];
122    while (!inflater.needsInput() || deflated.size() > 0 || deflatedIn.available() > 0) {
123      int count = inflatedIn.read(buffer, 0, buffer.length);
124      result.write(buffer, 0, count);
125    }
126    return result;
127  }
128
129  private ByteString randomBytes(int length) {
130    Random random = new Random(0);
131    byte[] randomBytes = new byte[length];
132    random.nextBytes(randomBytes);
133    return ByteString.of(randomBytes);
134  }
135
136  private String repeat(char c, int count) {
137    char[] array = new char[count];
138    Arrays.fill(array, c);
139    return new String(array);
140  }
141}
142