Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1739483 Fix S3 exception casting and flaky tests #1042

Merged
merged 6 commits into from
Oct 24, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
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
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
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
Loading