tests: fix flaky client_send_all case (#547)

Because the service returns early, it's possible the sink finish with errors. Signed-off-by: Jay Lee <[email protected]>
tikv · Sep 22, 2021 · 645c405 · 645c405
1 parent 7cd3b6b
commit 645c405
Showing 1 changed file with 18 additions and 9 deletions.
diff --git a/tests-and-examples/tests/cases/stream.rs b/tests-and-examples/tests/cases/stream.rs
@@ -49,9 +49,6 @@ impl RouteGuide for RouteGuideService {
                     break;
                 }
             }
-            // Sleep a while to avoid the possibility that the client will receive
-            // remotestopped error.
-            std::thread::sleep(std::time::Duration::from_millis(100));
             resp.success(summary).await?;
             Ok(())
         }
@@ -70,6 +67,16 @@ impl RouteGuide for RouteGuideService {
     }
 }
 
+macro_rules! assert_finish {
+    ($res:expr) => {
+        match $res {
+            // RouteGuide returns early, so `RpcFinished` is possible returned.
+            Ok(()) | Err(grpcio::Error::RpcFinished(None)) => (),
+            Err(e) => panic!("unexpected error {:?}", e),
+        }
+    };
+}
+
 #[test]
 fn test_client_send_all() {
     let env = Arc::new(EnvBuilder::new().build());
@@ -94,9 +101,10 @@ fn test_client_send_all() {
             send_data.push(p);
         }
         let send_stream = futures::stream::iter(send_data);
-        sink.send_all(&mut send_stream.map(move |item| Ok((item, WriteFlags::default()))))
-            .await
-            .unwrap();
+        assert_finish!(
+            sink.send_all(&mut send_stream.map(move |item| Ok((item, WriteFlags::default()))))
+                .await
+        );
         let summary = receiver.await.unwrap();
         assert_eq!(summary.get_point_count(), MESSAGE_NUM);
 
@@ -110,9 +118,10 @@ fn test_client_send_all() {
         }
         let send_stream = futures::stream::iter(send_data);
         sink.enhance_batch(true);
-        sink.send_all(&mut send_stream.map(move |item| Ok((item, WriteFlags::default()))))
-            .await
-            .unwrap();
+        assert_finish!(
+            sink.send_all(&mut send_stream.map(move |item| Ok((item, WriteFlags::default()))))
+                .await
+        );
         let summary = receiver.await.unwrap();
         assert_eq!(summary.get_point_count(), MESSAGE_NUM);