Skip to content

Commit

Permalink
SNOW-1739483 Fix S3 exception casting and flaky tests (#1042)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dstempniak
sfc-gh-dstempniak authored Oct 24, 2024
1 parent 9ae5119 commit 898a2f0
Show file tree
Hide file tree
Showing 5 changed files with 117 additions and 99 deletions.
8 changes: 4 additions & 4 deletions Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -514,8 +514,8 @@ public void TestDefaultLoginTimeout()

// Should timeout after the default timeout (300 sec)
Assert.GreaterOrEqual(stopwatch.ElapsedMilliseconds, conn.ConnectionTimeout * 1000 - delta);
// But never more because there's no connection timeout remaining
Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (conn.ConnectionTimeout + 1) * 1000);
// But never more because there's no connection timeout remaining (with 2 seconds margin)
Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (conn.ConnectionTimeout + 2) * 1000);
}
}
}
Expand Down Expand Up @@ -2015,8 +2015,8 @@ public void TestAsyncDefaultLoginTimeout()

// Should timeout after the default timeout (300 sec)
Assert.GreaterOrEqual(stopwatch.ElapsedMilliseconds, conn.ConnectionTimeout * 1000 - delta);
// But never more because there's no connection timeout remaining
Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (conn.ConnectionTimeout + 1) * 1000);
// But never more because there's no connection timeout remaining (with 2 seconds margin)
Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (conn.ConnectionTimeout + 2) * 1000);

Assert.AreEqual(ConnectionState.Closed, conn.State);
Assert.AreEqual(SFSessionHttpClientProperties.DefaultRetryTimeout.TotalSeconds, conn.ConnectionTimeout);
Expand Down
21 changes: 11 additions & 10 deletions Snowflake.Data.Tests/Mock/MockS3Client.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,22 +35,23 @@ class MockS3Client
internal const int ContentLength = 9999;

// Create AWS exception for mock requests
static Exception CreateMockAwsResponseError(string errorCode, bool isAsync)
static Exception CreateMockAwsResponseError(string awsErrorCode, bool isAsync)
{
AmazonS3Exception awsError = new AmazonS3Exception(S3ErrorMessage);
awsError.ErrorCode = errorCode;
Exception exception = awsErrorCode.Length > 0
? new AmazonS3Exception(S3ErrorMessage) { ErrorCode = awsErrorCode }
: new Exception("Non-AWS exception");

if (isAsync)
{
return awsError; // S3 throws the AmazonS3Exception on async calls
return exception; // S3 throws the AmazonS3Exception on async calls
}

Exception exceptionContainingS3Error = new Exception(S3ErrorMessage, awsError);
Exception exceptionContainingS3Error = new Exception(S3ErrorMessage, exception);
return exceptionContainingS3Error; // S3 places the AmazonS3Exception on the InnerException property on non-async calls
}

// Create mock response for GetFileHeader
static internal Task<GetObjectResponse> CreateResponseForGetFileHeader(string statusCode, bool isAsync)
internal static Task<GetObjectResponse> CreateResponseForGetFileHeader(string statusCode, bool isAsync)
{
if (statusCode == HttpStatusCode.OK.ToString())
{
Expand All @@ -70,20 +71,20 @@ static internal Task<GetObjectResponse> CreateResponseForGetFileHeader(string st
}

// Create mock response for UploadFile
static internal Task<PutObjectResponse> CreateResponseForUploadFile(string statusCode, bool isAsync)
internal static Task<PutObjectResponse> CreateResponseForUploadFile(string awsStatusCode, bool isAsync)
{
if (statusCode == HttpStatusCode.OK.ToString())
if (awsStatusCode == AwsStatusOk)
{
return Task.FromResult(new PutObjectResponse());
}
else
{
throw CreateMockAwsResponseError(statusCode, isAsync);
throw CreateMockAwsResponseError(awsStatusCode, isAsync);
}
}

// Create mock response for DownloadFile
static internal Task<GetObjectResponse> CreateResponseForDownloadFile(string statusCode, bool isAsync)
internal static Task<GetObjectResponse> CreateResponseForDownloadFile(string statusCode, bool isAsync)
{
if (statusCode == HttpStatusCode.OK.ToString())
{
Expand Down
54 changes: 18 additions & 36 deletions Snowflake.Data.Tests/UnitTests/SFS3ClientTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,18 +127,15 @@ public void TestExtractBucketNameAndPath()
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(SFS3Client.NO_SUCH_KEY, ResultStatus.NOT_FOUND_FILE)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.ERROR)] // Any error that isn't the above will return ResultStatus.ERROR
public void TestGetFileHeader(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.ERROR)] // For non-AWS exception will return ResultStatus.ERROR
public void TestGetFileHeader(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.GetObjectAsync(It.IsAny<GetObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<GetObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForGetFileHeader(request.BucketName, false);
});
.Returns(() => MockS3Client.CreateResponseForGetFileHeader(awsStatusCode, false));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;

// Act
FileHeader fileHeader = _client.GetFileHeader(_fileMetadata);
Expand All @@ -152,18 +149,15 @@ public void TestGetFileHeader(string requestKey, ResultStatus expectedResultStat
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(SFS3Client.NO_SUCH_KEY, ResultStatus.NOT_FOUND_FILE)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.ERROR)] // Any error that isn't the above will return ResultStatus.ERROR
public async Task TestGetFileHeaderAsync(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.ERROR)] // For non-AWS exception will return ResultStatus.ERROR
public async Task TestGetFileHeaderAsync(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.GetObjectAsync(It.IsAny<GetObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<GetObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForGetFileHeader(request.BucketName, true);
});
.Returns(() => MockS3Client.CreateResponseForGetFileHeader(awsStatusCode, true));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;

// Act
FileHeader fileHeader = await _client.GetFileHeaderAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false);
Expand Down Expand Up @@ -194,18 +188,15 @@ private void AssertForGetFileHeaderTests(ResultStatus expectedResultStatus, File
[TestCase(MockS3Client.AwsStatusOk, ResultStatus.UPLOADED)]
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.NEED_RETRY)] // Any error that isn't the above will return ResultStatus.NEED_RETRY
public void TestUploadFile(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.NEED_RETRY)] // For non-AWS exception will return ResultStatus.NEED_RETRY
public void TestUploadFile(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.PutObjectAsync(It.IsAny<PutObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<PutObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForUploadFile(request.BucketName, false);
});
.Returns(() => MockS3Client.CreateResponseForUploadFile(awsStatusCode, false));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;
_fileMetadata.uploadSize = UploadFileSize;

// Act
Expand Down Expand Up @@ -254,18 +245,15 @@ public void TestAppendHttpsToEndpointWithBrackets()
[TestCase(MockS3Client.AwsStatusOk, ResultStatus.UPLOADED)]
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.NEED_RETRY)] // Any error that isn't the above will return ResultStatus.NEED_RETRY
public async Task TestUploadFileAsync(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.NEED_RETRY)] // For non-AWS exception will return ResultStatus.NEED_RETRY
public async Task TestUploadFileAsync(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.PutObjectAsync(It.IsAny<PutObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<PutObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForUploadFile(request.BucketName, true);
});
.Returns(() => MockS3Client.CreateResponseForUploadFile(awsStatusCode, true));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;
_fileMetadata.uploadSize = UploadFileSize;

// Act
Expand Down Expand Up @@ -295,18 +283,15 @@ private void AssertForUploadFileTests(ResultStatus expectedResultStatus)
[TestCase(MockS3Client.AwsStatusOk, ResultStatus.DOWNLOADED)]
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.NEED_RETRY)] // Any error that isn't the above will return ResultStatus.NEED_RETRY
public void TestDownloadFile(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.NEED_RETRY)] // For non-AWS exception will return ResultStatus.NEED_RETRY
public void TestDownloadFile(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.GetObjectAsync(It.IsAny<GetObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<GetObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForDownloadFile(request.BucketName, false);
});
.Returns(() => MockS3Client.CreateResponseForDownloadFile(awsStatusCode, false));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;

// Act
_client.DownloadFile(_fileMetadata, t_downloadFileName, Parallel);
Expand All @@ -319,18 +304,15 @@ public void TestDownloadFile(string requestKey, ResultStatus expectedResultStatu
[TestCase(MockS3Client.AwsStatusOk, ResultStatus.DOWNLOADED)]
[TestCase(SFS3Client.EXPIRED_TOKEN, ResultStatus.RENEW_TOKEN)]
[TestCase(MockS3Client.AwsStatusError, ResultStatus.NEED_RETRY)] // Any error that isn't the above will return ResultStatus.NEED_RETRY
public async Task TestDownloadFileAsync(string requestKey, ResultStatus expectedResultStatus)
[TestCase("", ResultStatus.NEED_RETRY)] // For non-AWS exception will return ResultStatus.NEED_RETRY
public async Task TestDownloadFileAsync(string awsStatusCode, ResultStatus expectedResultStatus)
{
// Arrange
var mockAmazonS3Client = new Mock<AmazonS3Client>(AwsKeyId, AwsSecretKey, AwsToken, _clientConfig);
mockAmazonS3Client.Setup(client => client.GetObjectAsync(It.IsAny<GetObjectRequest>(), It.IsAny<CancellationToken>()))
.Returns<GetObjectRequest, CancellationToken>((request, cancellationToken) =>
{
return MockS3Client.CreateResponseForDownloadFile(request.BucketName, true);
});
.Returns(() => MockS3Client.CreateResponseForDownloadFile(awsStatusCode, true));
_client = new SFS3Client(_fileMetadata.stageInfo, MaxRetry, Parallel, _proxyCredentials, mockAmazonS3Client.Object);
_fileMetadata.client = _client;
_fileMetadata.stageInfo.location = requestKey;

// Act
await _client.DownloadFileAsync(_fileMetadata, t_downloadFileName, Parallel, _cancellationToken).ConfigureAwait(false);
Expand Down
Loading

0 comments on commit 898a2f0

Please sign in to comment.