From 58c4f0b5f41da115eda881fdfdba41a0e61f9854 Mon Sep 17 00:00:00 2001
From: Oleg Jukovec <oleg.jukovec@tarantool.org>
Date: Tue, 27 Dec 2022 12:54:37 +0300
Subject: [PATCH] bugfix: flaky TestClientRequestObjectsWithContext

The patch makes the test more deterministic. It helps to avoid
canceling a request with an already received response.

Closes #244
---
 CHANGELOG.md      |  2 ++
 tarantool_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7930b34d1..dab110777 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,8 @@ Versioning](http://semver.org/spec/v2.0.0.html) except to the first release.
 
 - Decimal package uses a test variable DecimalPrecision instead of a
   package-level variable decimalPrecision (#233)
+- Flaky tests TestClientRequestObjectsWithContext and
+  TestClientIdRequestObjectWithContext (#244)
 
 ## [1.9.0] - 2022-11-02
 
diff --git a/tarantool_test.go b/tarantool_test.go
index 9061db61c..82e531d37 100644
--- a/tarantool_test.go
+++ b/tarantool_test.go
@@ -2397,15 +2397,57 @@ func TestClientRequestObjectsWithPassedCanceledContext(t *testing.T) {
 	}
 }
 
+// waitCtxRequest waits for the WaitGroup in Body() call and returns
+// the context from Ctx() call. The request helps us to make sure that
+// the context's cancel() call is called before a response received.
+type waitCtxRequest struct {
+	ctx context.Context
+	wg  sync.WaitGroup
+}
+
+func (req *waitCtxRequest) Code() int32 {
+	return NewPingRequest().Code()
+}
+
+func (req *waitCtxRequest) Body(res SchemaResolver, enc *encoder) error {
+	req.wg.Wait()
+	return NewPingRequest().Body(res, enc)
+}
+
+func (req *waitCtxRequest) Ctx() context.Context {
+	return req.ctx
+}
+
+func (req *waitCtxRequest) Async() bool {
+	return NewPingRequest().Async()
+}
+
 func TestClientRequestObjectsWithContext(t *testing.T) {
 	var err error
 	conn := test_helpers.ConnectWithValidation(t, server, opts)
 	defer conn.Close()
 
 	ctx, cancel := context.WithCancel(context.Background())
-	req := NewPingRequest().Context(ctx)
-	fut := conn.Do(req)
+	req := &waitCtxRequest{ctx: ctx}
+	req.wg.Add(1)
+
+	var futWg sync.WaitGroup
+	var fut *Future
+
+	futWg.Add(1)
+	go func() {
+		defer futWg.Done()
+		fut = conn.Do(req)
+	}()
+
 	cancel()
+	req.wg.Done()
+
+	futWg.Wait()
+	if fut == nil {
+		t.Fatalf("fut must be not nil")
+	}
+
 	resp, err := fut.Get()
 	if resp != nil {
 		t.Fatalf("response must be nil")