test: add a simple test for atomic operations

Signed-off-by: ekexium <[email protected]>

Author: ekexium

mock-tikv/src/pd.rs

@@ -336,11 +336,11 @@ impl Pd for MockPd {
         todo!()
     }
 
-    fn get_dc_locations(
+    fn get_dc_location_info(
         &mut self,
-        _: ::grpcio::RpcContext<'_>,
-        _: GetDcLocationsRequest,
-        _: ::grpcio::UnarySink<GetDcLocationsResponse>,
+        _: grpcio::RpcContext,
+        _: GetDcLocationInfoRequest,
+        _: grpcio::UnarySink<GetDcLocationInfoResponse>,
     ) {
         todo!()
     }

mock-tikv/src/server.rs

@@ -5,7 +5,7 @@ use derive_new::new;
 use futures::{FutureExt, TryFutureExt};
 use grpcio::{Environment, Server, ServerBuilder};
 use std::sync::Arc;
-use tikv_client_proto::{kvrpcpb::*, tikvpb::*};
+use tikv_client_proto::{coprocessor_v2::*, kvrpcpb::*, tikvpb::*};
 
 pub const MOCK_TIKV_PORT: u16 = 50019;
 
@@ -513,4 +513,40 @@ impl Tikv for MockTikv {
     ) {
         todo!()
     }
+
+    fn raw_get_key_ttl(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: RawGetKeyTtlRequest,
+        _: grpcio::UnarySink<RawGetKeyTtlResponse>,
+    ) {
+        todo!()
+    }
+
+    fn raw_compare_and_swap(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: RawCasRequest,
+        _: grpcio::UnarySink<RawCasResponse>,
+    ) {
+        todo!()
+    }
+
+    fn coprocessor_v2(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: RawCoprocessorRequest,
+        _: grpcio::UnarySink<RawCoprocessorResponse>,
+    ) {
+        todo!()
+    }
+
+    fn get_store_safe_ts(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: StoreSafeTsRequest,
+        _: grpcio::UnarySink<StoreSafeTsResponse>,
+    ) {
+        todo!()
+    }
 }

src/raw/client.rs

@@ -449,10 +449,15 @@ impl Client {
     pub async fn atomic_compare_and_swap(
         &self,
         key: impl Into<Key>,
-        previous_value: Option<Value>,
-        new_value: Value,
+        previous_value: impl Into<Option<Value>>,
+        new_value: impl Into<Value>,
     ) -> Result<(Option<Value>, bool)> {
-        let req = new_cas_request(key.into(), new_value, previous_value, self.cf.clone());
+        let req = new_cas_request(
+            key.into(),
+            new_value.into(),
+            previous_value.into(),
+            self.cf.clone(),
+        );
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req)
             .single_region()
             .await?

tests/integration_tests.rs

@@ -748,6 +748,52 @@ async fn pessimistic_heartbeat() -> Result<()> {
 
     Ok(())
 }
+
+// It tests very basic functionality of atomic operations (put, cas, delete).
+#[tokio::test]
+#[serial]
+async fn raw_cas() -> Result<()> {
+    clear_tikv().await;
+    let client = RawClient::new(pd_addrs()).await?;
+    let key = "key".to_owned();
+    let value = "value".to_owned();
+    let new_value = "new value".to_owned();
+
+    client.atomic_put(key.clone(), value.clone()).await?;
+    assert_eq!(
+        client.get(key.clone()).await?.unwrap(),
+        value.clone().as_bytes()
+    );
+
+    client
+        .atomic_compare_and_swap(
+            key.clone(),
+            Some("another_value".to_owned()).map(|v| v.into()),
+            new_value.clone(),
+        )
+        .await?;
+    assert_ne!(
+        client.get(key.clone()).await?.unwrap(),
+        new_value.clone().as_bytes()
+    );
+
+    client
+        .atomic_compare_and_swap(
+            key.clone(),
+            Some(value.to_owned()).map(|v| v.into()),
+            new_value.clone(),
+        )
+        .await?;
+    assert_eq!(
+        client.get(key.clone()).await?.unwrap(),
+        new_value.clone().as_bytes()
+    );
+
+    client.atomic_delete(key.clone()).await?;
+    assert!(client.get(key.clone()).await?.is_none());
+    Ok(())
+}
+
 // helper function
 async fn get_u32(client: &RawClient, key: Vec<u8>) -> Result<u32> {
     let x = client.get(key).await?.unwrap();

tikv-client-proto/src/lib.rs

@@ -1,7 +1,9 @@
 // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
 
 use protos::*;
-pub use protos::{coprocessor, errorpb, kvrpcpb, metapb, mpp, pdpb, raft_serverpb, tikvpb};
+pub use protos::{
+    coprocessor, coprocessor_v2, errorpb, kvrpcpb, metapb, mpp, pdpb, raft_serverpb, tikvpb,
+};
 
 #[allow(dead_code)]
 #[allow(clippy::all)]