From b1efff659e82a1f137a4b70fcd2a67c09237cef5 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Thu, 12 Mar 2020 11:24:36 +0100
Subject: [PATCH] rpc: improve cancel test (#20752)

This is supposed to fix the occasional failures in
TestCancel* on Travis CI.
---
 rpc/client_test.go      | 16 +++++++++-------
 rpc/server_test.go      |  2 +-
 rpc/testservice_test.go |  5 +++++
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/rpc/client_test.go b/rpc/client_test.go
index 0287dad088..1d87d96671 100644
--- a/rpc/client_test.go
+++ b/rpc/client_test.go
@@ -179,7 +179,7 @@ func testClientCancel(transport string, t *testing.T) {
 	var (
 		wg       sync.WaitGroup
 		nreqs    = 10
-		ncallers = 6
+		ncallers = 10
 	)
 	caller := func(index int) {
 		defer wg.Done()
@@ -200,14 +200,16 @@ func testClientCancel(transport string, t *testing.T) {
 				// deadline.
 				ctx, cancel = context.WithTimeout(context.Background(), timeout)
 			}
+
 			// Now perform a call with the context.
 			// The key thing here is that no call will ever complete successfully.
-			sleepTime := maxContextCancelTimeout + 20*time.Millisecond
-			err := client.CallContext(ctx, nil, "test_sleep", sleepTime)
-			if err != nil {
-				log.Debug(fmt.Sprint("got expected error:", err))
-			} else {
-				t.Errorf("no error for call with %v wait time", timeout)
+			err := client.CallContext(ctx, nil, "test_block")
+			switch {
+			case err == nil:
+				_, hasDeadline := ctx.Deadline()
+				t.Errorf("no error for call with %v wait time (deadline: %v)", timeout, hasDeadline)
+				// default:
+				// 	t.Logf("got expected error with %v wait time: %v", timeout, err)
 			}
 			cancel()
 		}
diff --git a/rpc/server_test.go b/rpc/server_test.go
index a4ca1fde43..99cca26ddf 100644
--- a/rpc/server_test.go
+++ b/rpc/server_test.go
@@ -45,7 +45,7 @@ func TestServerRegisterName(t *testing.T) {
 		t.Fatalf("Expected service calc to be registered")
 	}
 
-	wantCallbacks := 7
+	wantCallbacks := 8
 	if len(svc.callbacks) != wantCallbacks {
 		t.Errorf("Expected %d callbacks for service 'service', got %d", wantCallbacks, len(svc.callbacks))
 	}
diff --git a/rpc/testservice_test.go b/rpc/testservice_test.go
index 010fc4f0bb..3094c639e6 100644
--- a/rpc/testservice_test.go
+++ b/rpc/testservice_test.go
@@ -77,6 +77,11 @@ func (s *testService) Sleep(ctx context.Context, duration time.Duration) {
 	time.Sleep(duration)
 }
 
+func (s *testService) Block(ctx context.Context) error {
+	<-ctx.Done()
+	return errors.New("context canceled in testservice_block")
+}
+
 func (s *testService) Rets() (string, error) {
 	return "", nil
 }