Add simple tests for zstd streams

src/test/java/com/google/devtools/build/lib/remote/BUILD

@@ -16,6 +16,7 @@ filegroup(
         "//src/test/java/com/google/devtools/build/lib/remote/merkletree:srcs",
         "//src/test/java/com/google/devtools/build/lib/remote/options:srcs",
         "//src/test/java/com/google/devtools/build/lib/remote/util:srcs",
+        "//src/test/java/com/google/devtools/build/lib/remote/zstd:srcs",
     ],
     visibility = ["//src/test/java/com/google/devtools/build/lib:__pkg__"],
 )

src/test/java/com/google/devtools/build/lib/remote/zstd/BUILD

@@ -0,0 +1,28 @@
+load("@rules_java//java:defs.bzl", "java_test")
+
+package(
+    default_testonly = 1,
+    default_visibility = ["//src:__subpackages__"],
+)
+
+licenses(["notice"])
+
+filegroup(
+    name = "srcs",
+    testonly = 0,
+    srcs = glob(["**"]),
+    visibility = ["//src:__subpackages__"],
+)
+
+java_test(
+    name = "zstd",
+    srcs = glob(["*.java"]),
+    test_class = "com.google.devtools.build.lib.AllTests",
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/remote/zstd:zstd",
+        "//src/test/java/com/google/devtools/build/lib:test_runner",
+        "//third_party:junit4",
+        "//third_party:truth",
+        "@zstd-jni//:zstd-jni",
+    ],
+)

src/test/java/com/google/devtools/build/lib/remote/zstd/ZstdCompressingInputStreamTest.java

@@ -0,0 +1,48 @@
+package com.google.devtools.build.lib.remote.zstd;
+
+import com.github.luben.zstd.Zstd;
+import com.google.common.io.ByteStreams;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.Random;
+
+import static com.google.common.truth.Truth.assertThat;
+
+@RunWith(JUnit4.class)
+public class ZstdCompressingInputStreamTest {
+  @Test
+  public void compressionWorks() throws IOException {
+    Random rand = new Random();
+    byte[] data = new byte[50];
+    rand.nextBytes(data);
+
+    ByteArrayInputStream bais = new ByteArrayInputStream(data);
+    ZstdCompressingInputStream zdis = new ZstdCompressingInputStream(bais);
+
+    ByteArrayOutputStream compressed = new ByteArrayOutputStream();
+    ByteStreams.copy(zdis, compressed);
+
+    assertThat(Zstd.decompress(compressed.toByteArray(), data.length)).isEqualTo(data);
+  }
+
+  @Test
+  public void streamCanBeCompressedWithMinimumBufferSize() throws IOException {
+    Random rand = new Random();
+    byte[] data = new byte[50];
+    rand.nextBytes(data);
+
+    ByteArrayInputStream bais = new ByteArrayInputStream(data);
+    ZstdCompressingInputStream zdis =
+        new ZstdCompressingInputStream(bais, ZstdCompressingInputStream.MIN_BUFFER_SIZE);
+
+    ByteArrayOutputStream compressed = new ByteArrayOutputStream();
+    ByteStreams.copy(zdis, compressed);
+
+    assertThat(Zstd.decompress(compressed.toByteArray(), data.length)).isEqualTo(data);
+  }
+}

src/test/java/com/google/devtools/build/lib/remote/zstd/ZstdDecompressingOutputStreamTest.java

@@ -0,0 +1,74 @@
+package com.google.devtools.build.lib.remote.zstd;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.Random;
+
+@RunWith(JUnit4.class)
+public class ZstdDecompressingOutputStreamTest {
+  @Test
+  public void decompressionWorks() throws IOException {
+    Random rand = new Random();
+    byte[] data = new byte[50];
+    rand.nextBytes(data);
+    byte[] compressed = Zstd.compress(data);
+
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    ZstdDecompressingOutputStream zdos = new ZstdDecompressingOutputStream(baos);
+    zdos.write(compressed);
+    zdos.flush();
+
+    assertThat(baos.toByteArray()).isEqualTo(data);
+  }
+
+  @Test
+  public void streamCanBeDecompressedOneByteAtATime() throws IOException {
+    Random rand = new Random();
+    byte[] data = new byte[50];
+    rand.nextBytes(data);
+    byte[] compressed = Zstd.compress(data);
+
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    ZstdDecompressingOutputStream zdos = new ZstdDecompressingOutputStream(baos);
+    for (byte b : compressed) {
+      zdos.write(b);
+    }
+    zdos.flush();
+
+    assertThat(baos.toByteArray()).isEqualTo(data);
+  }
+
+  @Test
+  public void bytesWrittenMatchesDecompressedBytes() throws IOException {
+    byte[] data = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA".getBytes();
+
+    ByteArrayOutputStream compressed = new ByteArrayOutputStream();
+    ZstdOutputStream zos = new ZstdOutputStream(compressed);
+    zos.setCloseFrameOnFlush(true);
+    for (int i = 0; i < data.length; i++) {
+      zos.write(data[i]);
+      if (i % 5 == 0) {
+        // Create multiple frames of 5 bytes each.
+        zos.flush();
+      }
+    }
+    zos.close();
+
+    ByteArrayOutputStream decompressed = new ByteArrayOutputStream();
+    ZstdDecompressingOutputStream zdos = new ZstdDecompressingOutputStream(decompressed);
+    for (byte b : compressed.toByteArray()) {
+      zdos.write(b);
+      zdos.flush();
+      assertThat(zdos.getBytesWritten()).isEqualTo(decompressed.toByteArray().length);
+    }
+    assertThat(decompressed.toByteArray()).isEqualTo(data);
+  }
+}