From 3af031b434e65ee19b71048603c59bacc3d159a4 Mon Sep 17 00:00:00 2001
From: neverchanje <wutao1@xiaomi.com>
Date: Wed, 21 Nov 2018 10:46:55 +0800
Subject: [PATCH] add unit test

---
 .../infra/pegasus/tools/ZstdWrapper.java      |  3 +-
 .../infra/pegasus/tools/TestZstdWrapper.java  | 34 ++++++++++++++++---
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/main/java/com/xiaomi/infra/pegasus/tools/ZstdWrapper.java b/src/main/java/com/xiaomi/infra/pegasus/tools/ZstdWrapper.java
index 80852a21..4ce58ce7 100644
--- a/src/main/java/com/xiaomi/infra/pegasus/tools/ZstdWrapper.java
+++ b/src/main/java/com/xiaomi/infra/pegasus/tools/ZstdWrapper.java
@@ -77,6 +77,7 @@ public static byte[] decompress(byte[] src) throws PException {
 
     /**
      * A handle for Zstd dictionary.
+     * Dict can be created once and shared by multiple threads concurrently, since its usage is read-only.
      */
     public static class Dict {
         private Dict(byte[] dictBuffer) {
@@ -119,7 +120,7 @@ public void trainAndSaveResultToPegasus(PegasusTableInterface table) throws PExc
         }
 
         private void train() throws PException {
-            if (samples.size() < 100) {
+            if (samples.size() < 200) {
                 throw new PException("samples are not enough for training");
             }
 
diff --git a/src/test/java/com/xiaomi/infra/pegasus/tools/TestZstdWrapper.java b/src/test/java/com/xiaomi/infra/pegasus/tools/TestZstdWrapper.java
index a167f35d..e6280a6f 100644
--- a/src/test/java/com/xiaomi/infra/pegasus/tools/TestZstdWrapper.java
+++ b/src/test/java/com/xiaomi/infra/pegasus/tools/TestZstdWrapper.java
@@ -1,10 +1,7 @@
 package com.xiaomi.infra.pegasus.tools;
 
-
-import com.xiaomi.infra.pegasus.client.PException;
-import com.xiaomi.infra.pegasus.client.PegasusClientFactory;
-import com.xiaomi.infra.pegasus.client.PegasusClientInterface;
-import com.xiaomi.infra.pegasus.client.PegasusTableInterface;
+import com.xiaomi.infra.pegasus.client.*;
+import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -57,4 +54,31 @@ public void testCompression() throws Exception {
             }
         }
     }
+
+    @Test
+    public void testCompressionUsingDict() throws Exception {
+        PegasusClientInterface client = PegasusClientFactory.getSingletonClient();
+        PegasusTableInterface table = client.openTable("wutao1_sample");
+
+        ZstdWrapper.DictTrainer trainer = ZstdWrapper.startTraining();
+        for (int i = 0; i < 200; i++) {
+            byte[] value = RandomStringUtils.randomAlphabetic(1000).getBytes();
+            trainer.addSample(value);
+        }
+
+        // start training
+        trainer.trainAndSaveResultToPegasus(table);
+
+        // get result
+        ZstdWrapper.Dict dict = ZstdWrapper.getDictFromPegasus(table);
+
+        for (int i = 0; i < 10000; i++) {
+            byte[] value = RandomStringUtils.randomAlphabetic(1000).getBytes();
+
+            byte[] compressed = ZstdWrapper.compress(value, dict);
+            byte[] original = ZstdWrapper.decompress(compressed, dict);
+
+            Assert.assertArrayEquals(value, original);
+        }
+    }
 }