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

@@ -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);
     }
 

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()
     }

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");