Skip to content

Commit 308dbf2

Browse files
authored
Fix flaky e2e tests for gae and cloud run (#87)
* add unit tests for testclient.go * Add multiplexing to the test client to fix flakes
1 parent 82ad8f7 commit 308dbf2

3 files changed

Lines changed: 190 additions & 31 deletions

File tree

e2etestrunner/testclient/testclient.go

Lines changed: 59 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"github.com/GoogleCloudPlatform/opentelemetry-operations-e2e-testing/e2etesting/setuptf"
2121
"log"
2222
"strconv"
23+
"sync"
2324

2425
"cloud.google.com/go/pubsub"
2526
"google.golang.org/genproto/googleapis/rpc/code"
@@ -48,6 +49,14 @@ type Client struct {
4849
pubsubClient *pubsub.Client
4950
requestTopic *pubsub.Topic
5051
responseSubscription *pubsub.Subscription
52+
53+
mu sync.Mutex
54+
pendingRequests map[string]chan asyncResponse
55+
}
56+
57+
type asyncResponse struct {
58+
res *Response
59+
err error
5160
}
5261

5362
func New(ctx context.Context, projectID string, pubsubInfo *setuptf.PubsubInfo) (*Client, error) {
@@ -59,12 +68,44 @@ func New(ctx context.Context, projectID string, pubsubInfo *setuptf.PubsubInfo)
5968
pubsubClient: pubsub,
6069
requestTopic: pubsub.Topic(pubsubInfo.RequestTopic.TopicName),
6170
responseSubscription: pubsub.Subscription(pubsubInfo.ResponseTopic.SubscriptionName),
71+
pendingRequests: make(map[string]chan asyncResponse),
6272
}
6373
// Disable buffering
6474
client.requestTopic.PublishSettings.CountThreshold = 1
75+
76+
go client.startReceiver(ctx)
77+
6578
return client, nil
6679
}
6780

81+
func (c *Client) startReceiver(ctx context.Context) {
82+
err := c.responseSubscription.Receive(ctx, func(ctx context.Context, message *pubsub.Message) {
83+
testID := message.Attributes[TestID]
84+
85+
c.mu.Lock()
86+
ch, ok := c.pendingRequests[testID]
87+
c.mu.Unlock()
88+
89+
if ok {
90+
message.Ack()
91+
var resErr error
92+
var res *Response
93+
codeInt, err := strconv.Atoi(message.Attributes[StatusCode])
94+
if err != nil {
95+
resErr = fmt.Errorf(`response pub/sub message invalid attribute %q: %v, message: %v`, StatusCode, err, message)
96+
} else {
97+
res = &Response{StatusCode: code.Code(codeInt), Headers: message.Attributes}
98+
}
99+
ch <- asyncResponse{res: res, err: resErr}
100+
} else {
101+
message.Nack()
102+
}
103+
})
104+
if err != nil {
105+
log.Printf("Background subscriber error: %v", err)
106+
}
107+
}
108+
68109
func (c *Client) Request(
69110
ctx context.Context,
70111
request Request,
@@ -73,6 +114,18 @@ func (c *Client) Request(
73114
for k, v := range request.Headers {
74115
attributes[k] = v
75116
}
117+
118+
resCh := make(chan asyncResponse, 1)
119+
c.mu.Lock()
120+
c.pendingRequests[request.TestID] = resCh
121+
c.mu.Unlock()
122+
123+
defer func() {
124+
c.mu.Lock()
125+
delete(c.pendingRequests, request.TestID)
126+
c.mu.Unlock()
127+
}()
128+
76129
pubResult := c.requestTopic.Publish(ctx, &pubsub.Message{
77130
Attributes: attributes,
78131
})
@@ -81,40 +134,17 @@ func (c *Client) Request(
81134
return nil, err
82135
}
83136

84-
var (
85-
res *Response
86-
resErr error
87-
)
88-
cctx, cancel := context.WithCancel(ctx)
89-
err = c.responseSubscription.Receive(cctx, func(ctx context.Context, message *pubsub.Message) {
90-
if testID := message.Attributes[TestID]; testID == request.TestID {
91-
message.Ack()
92-
codeInt, err := strconv.Atoi(message.Attributes[StatusCode])
93-
if err != nil {
94-
resErr = fmt.Errorf(`response pub/sub message invalid attribute %q: %v, message: %v`, StatusCode, err, message)
95-
} else {
96-
res = &Response{StatusCode: code.Code(codeInt), Headers: message.Attributes}
97-
}
98-
cancel()
99-
} else {
100-
message.Nack()
101-
}
102-
})
103-
104-
if err != nil {
105-
return nil, err
106-
} else if resErr != nil {
107-
return nil, resErr
108-
} else if res == nil {
109-
// Can happen if cctx times out
137+
select {
138+
case <-ctx.Done():
110139
return nil, fmt.Errorf(
111-
"sent message ID %v, but never received a response on subscription %v",
140+
"sent message ID %v, but never received a response on subscription %v: %w",
112141
messageID,
113142
c.responseSubscription.String(),
143+
ctx.Err(),
114144
)
145+
case asyncRes := <-resCh:
146+
return asyncRes.res, asyncRes.err
115147
}
116-
117-
return res, nil
118148
}
119149

120150
// Call in TestMain() to block until the test server is ready for requests. Uses
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
package testclient
2+
3+
import (
4+
"context"
5+
"os"
6+
"strconv"
7+
"sync"
8+
"testing"
9+
10+
"cloud.google.com/go/pubsub"
11+
"cloud.google.com/go/pubsub/pstest"
12+
"github.com/GoogleCloudPlatform/opentelemetry-operations-e2e-testing/e2etesting/setuptf"
13+
"github.com/stretchr/testify/assert"
14+
"github.com/stretchr/testify/require"
15+
"google.golang.org/api/option"
16+
"google.golang.org/genproto/googleapis/rpc/code"
17+
"google.golang.org/grpc"
18+
"google.golang.org/grpc/credentials/insecure"
19+
)
20+
21+
func TestClientRequest(t *testing.T) {
22+
ctx := context.Background()
23+
24+
srv := pstest.NewServer()
25+
defer srv.Close()
26+
27+
os.Setenv("PUBSUB_EMULATOR_HOST", srv.Addr)
28+
defer os.Unsetenv("PUBSUB_EMULATOR_HOST")
29+
30+
conn, err := grpc.Dial(srv.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
31+
require.NoError(t, err)
32+
defer conn.Close()
33+
34+
client, err := pubsub.NewClient(ctx, "project", option.WithGRPCConn(conn))
35+
require.NoError(t, err)
36+
defer client.Close()
37+
38+
pubsubInfo := &setuptf.PubsubInfo{
39+
RequestTopic: setuptf.TopicInfo{
40+
TopicName: "request-topic",
41+
},
42+
ResponseTopic: setuptf.TopicInfo{
43+
TopicName: "response-topic",
44+
SubscriptionName: "response-sub",
45+
},
46+
}
47+
48+
reqTopic, err := client.CreateTopic(ctx, "request-topic")
49+
require.NoError(t, err)
50+
reqSub, err := client.CreateSubscription(ctx, "request-sub", pubsub.SubscriptionConfig{
51+
Topic: reqTopic,
52+
})
53+
require.NoError(t, err)
54+
55+
respTopic, err := client.CreateTopic(ctx, "response-topic")
56+
require.NoError(t, err)
57+
58+
_, err = client.CreateSubscription(ctx, "response-sub", pubsub.SubscriptionConfig{
59+
Topic: respTopic,
60+
})
61+
require.NoError(t, err)
62+
63+
// Mock server answering requests
64+
go func() {
65+
err := reqSub.Receive(ctx, func(c context.Context, msg *pubsub.Message) {
66+
msg.Ack()
67+
68+
custom := "header"
69+
if val, ok := msg.Attributes["foo"]; ok {
70+
custom = val
71+
}
72+
73+
// Send response
74+
res := &pubsub.Message{
75+
Attributes: map[string]string{
76+
TestID: msg.Attributes[TestID],
77+
StatusCode: strconv.Itoa(int(code.Code_OK)),
78+
"custom": custom,
79+
},
80+
}
81+
respTopic.Publish(c, res)
82+
})
83+
if err != nil {
84+
t.Logf("reqSub receive error: %v", err)
85+
}
86+
}()
87+
88+
sut, err := New(ctx, "project", pubsubInfo)
89+
require.NoError(t, err)
90+
91+
t.Run("single request", func(t *testing.T) {
92+
req := Request{
93+
TestID: "test-123",
94+
Scenario: "my-scenario",
95+
Headers: map[string]string{"foo": "bar"},
96+
}
97+
98+
resp, err := sut.Request(ctx, req)
99+
require.NoError(t, err)
100+
require.NotNil(t, resp)
101+
assert.Equal(t, code.Code_OK, resp.StatusCode)
102+
assert.Equal(t, "bar", resp.Headers["custom"])
103+
assert.Equal(t, "test-123", resp.Headers[TestID])
104+
})
105+
106+
t.Run("multiplexing", func(t *testing.T) {
107+
var wg sync.WaitGroup
108+
for i := 0; i < 10; i++ {
109+
wg.Add(1)
110+
go func(idx int) {
111+
defer wg.Done()
112+
req := Request{
113+
TestID: "test-" + strconv.Itoa(idx),
114+
Scenario: "my-scenario-" + strconv.Itoa(idx),
115+
Headers: map[string]string{"foo": "bar-" + strconv.Itoa(idx)},
116+
}
117+
118+
resp, err := sut.Request(ctx, req)
119+
assert.NoError(t, err)
120+
require.NotNil(t, resp)
121+
assert.Equal(t, code.Code_OK, resp.StatusCode)
122+
assert.Equal(t, "bar-"+strconv.Itoa(idx), resp.Headers["custom"])
123+
assert.Equal(t, "test-"+strconv.Itoa(idx), resp.Headers[TestID])
124+
}(i)
125+
}
126+
wg.Wait()
127+
})
128+
}

go.mod

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ require (
2929
golang.org/x/sync v0.10.0
3030
google.golang.org/api v0.196.0
3131
google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1
32+
google.golang.org/grpc v1.66.0
3233
google.golang.org/protobuf v1.34.2
3334
)
3435

@@ -50,6 +51,7 @@ require (
5051
github.com/go-logr/stdr v1.2.2 // indirect
5152
github.com/gogo/protobuf v1.3.2 // indirect
5253
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
54+
github.com/google/go-cmp v0.6.0 // indirect
5355
github.com/google/s2a-go v0.1.8 // indirect
5456
github.com/google/uuid v1.6.0 // indirect
5557
github.com/googleapis/enterprise-certificate-proxy v0.3.3 // indirect
@@ -63,6 +65,7 @@ require (
6365
github.com/pkg/errors v0.9.1 // indirect
6466
github.com/pmezard/go-difflib v1.0.0 // indirect
6567
github.com/rogpeppe/go-internal v1.12.0 // indirect
68+
go.einride.tech/aip v0.67.1 // indirect
6669
go.opencensus.io v0.24.0 // indirect
6770
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
6871
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
@@ -79,8 +82,6 @@ require (
7982
golang.org/x/time v0.6.0 // indirect
8083
google.golang.org/genproto v0.0.0-20240903143218-8af14fe29dc1 // indirect
8184
google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1 // indirect
82-
google.golang.org/grpc v1.66.0 // indirect
8385
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
8486
gopkg.in/yaml.v3 v3.0.1 // indirect
85-
gotest.tools/v3 v3.5.1 // indirect
8687
)

0 commit comments

Comments
 (0)