Add unit tests for SQS rendezvous in client

Co-authored-by: Michael Pu <michael.pu@uwaterloo.ca>
This commit is contained in:
Anthony Chang 2024-01-12 23:23:33 -05:00 committed by Cecylia Bocovich
parent f3b062ddb2
commit 32e864b71d
No known key found for this signature in database
GPG key ID: 009DE379FD9B7B90
3 changed files with 137 additions and 38 deletions

View file

@ -7,12 +7,18 @@ import (
"io"
"io/ioutil"
"net/http"
"net/url"
"testing"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/sqs"
"github.com/aws/aws-sdk-go-v2/service/sqs/types"
"github.com/golang/mock/gomock"
. "github.com/smartystreets/goconvey/convey"
"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/amp"
"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/messages"
"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/nat"
"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/sqsclient"
)
// mockTransport's RoundTrip method returns a response with a fake status and
@ -271,3 +277,123 @@ func TestAMPCacheRendezvous(t *testing.T) {
})
})
}
func TestSQSRendezvous(t *testing.T) {
Convey("SQS Rendezvous", t, func() {
Convey("Construct SQS queue rendezvous", func() {
transport := &mockTransport{http.StatusOK, []byte{}}
rend, err := newSQSRendezvous("https://sqs.us-east-1.amazonaws.com", "some-access-key-id", "some-secret-key", transport)
So(err, ShouldBeNil)
So(rend.sqsClientID, ShouldNotBeNil)
So(rend.sqsClient, ShouldNotBeNil)
So(rend.sqsURL, ShouldNotBeNil)
So(rend.sqsURL.String(), ShouldResemble, "https://sqs.us-east-1.amazonaws.com")
})
ctrl := gomock.NewController(t)
mockSqsClient := sqsclient.NewMockSQSClient(ctrl)
responseQueueURL := "https://sqs.us-east-1.amazonaws.com/testing"
sqsClientID := "test123"
sqsUrl, _ := url.Parse("https://sqs.us-east-1.amazonaws.com/broker")
fakeEncPollResp := makeEncPollResp(
`{"answer": "{\"type\":\"answer\",\"sdp\":\"fake\"}" }`,
"",
)
sqsRendezvous := sqsRendezvous{
transport: &mockTransport{http.StatusOK, []byte{}},
sqsClientID: sqsClientID,
sqsClient: mockSqsClient,
sqsURL: sqsUrl,
timeout: 0,
numRetries: 5,
}
Convey("sqsRendezvous.Exchange responds with answer", func() {
mockSqsClient.EXPECT().SendMessage(gomock.Any(), &sqs.SendMessageInput{
MessageAttributes: map[string]types.MessageAttributeValue{
"ClientID": {
DataType: aws.String("String"),
StringValue: aws.String(sqsClientID),
},
},
MessageBody: aws.String(string(fakeEncPollResp)),
QueueUrl: aws.String(sqsUrl.String()),
})
mockSqsClient.EXPECT().GetQueueUrl(gomock.Any(), &sqs.GetQueueUrlInput{
QueueName: aws.String("snowflake-client-" + sqsClientID),
}).Return(&sqs.GetQueueUrlOutput{
QueueUrl: aws.String(responseQueueURL),
}, nil)
mockSqsClient.EXPECT().ReceiveMessage(gomock.Any(), gomock.Eq(&sqs.ReceiveMessageInput{
QueueUrl: &responseQueueURL,
MaxNumberOfMessages: 1,
WaitTimeSeconds: 20,
})).Return(&sqs.ReceiveMessageOutput{
Messages: []types.Message{{Body: aws.String("answer")}},
}, nil)
answer, err := sqsRendezvous.Exchange(fakeEncPollResp)
So(answer, ShouldEqual, []byte("answer"))
So(err, ShouldBeNil)
})
Convey("sqsRendezvous.Exchange cannot get queue url", func() {
mockSqsClient.EXPECT().SendMessage(gomock.Any(), &sqs.SendMessageInput{
MessageAttributes: map[string]types.MessageAttributeValue{
"ClientID": {
DataType: aws.String("String"),
StringValue: aws.String(sqsClientID),
},
},
MessageBody: aws.String(string(fakeEncPollResp)),
QueueUrl: aws.String(sqsUrl.String()),
})
for i := 0; i < sqsRendezvous.numRetries; i++ {
mockSqsClient.EXPECT().GetQueueUrl(gomock.Any(), &sqs.GetQueueUrlInput{
QueueName: aws.String("snowflake-client-" + sqsClientID),
}).Return(nil, errors.New("test error"))
}
answer, err := sqsRendezvous.Exchange(fakeEncPollResp)
So(answer, ShouldBeNil)
So(err, ShouldNotBeNil)
So(err, ShouldEqual, errors.New("test error"))
})
Convey("sqsRendezvous.Exchange does not receive answer", func() {
mockSqsClient.EXPECT().SendMessage(gomock.Any(), &sqs.SendMessageInput{
MessageAttributes: map[string]types.MessageAttributeValue{
"ClientID": {
DataType: aws.String("String"),
StringValue: aws.String(sqsClientID),
},
},
MessageBody: aws.String(string(fakeEncPollResp)),
QueueUrl: aws.String(sqsUrl.String()),
})
mockSqsClient.EXPECT().GetQueueUrl(gomock.Any(), &sqs.GetQueueUrlInput{
QueueName: aws.String("snowflake-client-" + sqsClientID),
}).Return(&sqs.GetQueueUrlOutput{
QueueUrl: aws.String(responseQueueURL),
}, nil)
for i := 0; i < sqsRendezvous.numRetries; i++ {
mockSqsClient.EXPECT().ReceiveMessage(gomock.Any(), gomock.Eq(&sqs.ReceiveMessageInput{
QueueUrl: &responseQueueURL,
MaxNumberOfMessages: 1,
WaitTimeSeconds: 20,
})).Return(&sqs.ReceiveMessageOutput{
Messages: []types.Message{},
}, nil)
}
answer, err := sqsRendezvous.Exchange(fakeEncPollResp)
So(answer, ShouldEqual, []byte{})
So(err, ShouldBeNil)
})
})
}