From 1bf9374939e8f61865f26262cbc1ca7f6458ed3c Mon Sep 17 00:00:00 2001
From: ekexium <ekexium@gmail.com>
Date: Mon, 14 Sep 2020 12:17:24 +0800
Subject: [PATCH] mock tikv: distinguish between empty value and non-exsit
 keys; same behavior with TiKV

Signed-off-by: ekexium <ekexium@gmail.com>
---
 mock-tikv/src/server.rs  |  6 +++++-
 mock-tikv/src/store.rs   | 18 +++++++++++-------
 tests/mock_tikv_tests.rs | 10 ++++------
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/mock-tikv/src/server.rs b/mock-tikv/src/server.rs
index abacbf5..b3b5e9a 100644
--- a/mock-tikv/src/server.rs
+++ b/mock-tikv/src/server.rs
@@ -178,7 +178,11 @@ impl Tikv for MockTikv {
         sink: grpcio::UnarySink<kvproto::kvrpcpb::RawGetResponse>,
     ) {
         let mut resp = RawGetResponse::default();
-        resp.set_value(self.inner.raw_get(req.get_key()));
+        if let Some(v) = self.inner.raw_get(req.get_key()) {
+            resp.set_value(v);
+        } else {
+            resp.set_not_found(true);
+        }
         spawn_unary_success!(ctx, req, resp, sink);
     }
 
diff --git a/mock-tikv/src/store.rs b/mock-tikv/src/store.rs
index 80c95d7..70c5f52 100644
--- a/mock-tikv/src/store.rs
+++ b/mock-tikv/src/store.rs
@@ -24,19 +24,23 @@ impl KvStore {
         }
     }
 
-    pub fn raw_get(&self, key: &[u8]) -> Vec<u8> {
+    pub fn raw_get(&self, key: &[u8]) -> Option<Vec<u8>> {
         let data = self.data.read().unwrap();
-        data.get(key).map(|v| v.to_vec()).unwrap_or_else(Vec::new)
+        data.get(key).map(|v| v.to_vec())
     }
 
     pub fn raw_batch_get(&self, keys: &[Vec<u8>]) -> Vec<KvPair> {
         let data = self.data.read().unwrap();
         keys.iter()
-            .map(|key| {
-                let mut pair = KvPair::default();
-                pair.set_value(data.get(key).map(|v| v.to_vec()).unwrap_or_else(Vec::new));
-                pair.set_key(key.to_vec());
-                pair
+            .filter_map(|key| {
+                if data.contains_key(key) {
+                    let mut pair = KvPair::default();
+                    pair.set_value(data.get(key).map(|v| v.to_vec()).unwrap_or_else(Vec::new));
+                    pair.set_key(key.to_vec());
+                    Some(pair)
+                } else {
+                    None
+                }
             })
             .collect()
     }
diff --git a/tests/mock_tikv_tests.rs b/tests/mock_tikv_tests.rs
index 840ff80..2bc07e6 100644
--- a/tests/mock_tikv_tests.rs
+++ b/tests/mock_tikv_tests.rs
@@ -20,7 +20,7 @@ mod test {
 
         // empty; get non-existent key
         let res = client.get("k1".to_owned()).await;
-        assert_eq!(res.unwrap().unwrap(), vec![]);
+        assert_eq!(res.unwrap(), None);
 
         // empty; put then batch_get
         let _ = client.put("k1".to_owned(), "v1".to_owned()).await.unwrap();
@@ -30,9 +30,9 @@ mod test {
             .batch_get(vec!["k1".to_owned(), "k2".to_owned(), "k3".to_owned()])
             .await
             .unwrap();
+        assert_eq!(res.len(), 2);
         assert_eq!(res[0].1, "v1".as_bytes());
         assert_eq!(res[1].1, "v2".as_bytes());
-        assert_eq!(res[2].1, "".as_bytes());
 
         // k1,k2; batch_put then batch_get
         let _ = client
@@ -58,7 +58,7 @@ mod test {
         assert!(res.is_err());
 
         let res = client.get("k3".to_owned()).await;
-        assert_eq!(res.unwrap().unwrap(), "".as_bytes());
+        assert_eq!(res.unwrap(), None);
 
         // k1,k2,k4; batch_delete then batch_get
         let res = client
@@ -75,9 +75,7 @@ mod test {
             ])
             .await
             .unwrap();
-        for i in 0..3 {
-            assert_eq!(res[i].1, "".as_bytes());
-        }
+        assert_eq!(res.len(), 0);
 
         debug!("Pass all tests");