From bd89073989ac5d5e8d8be9f58449339d942cc9ab Mon Sep 17 00:00:00 2001 From: sfc-gh-ext-simba-lf <115584722+sfc-gh-ext-simba-lf@users.noreply.github.com> Date: Wed, 20 Sep 2023 09:02:23 -0700 Subject: [PATCH] SNOW-723810: Add coverage for SFRemoteStorageUtil (#773) ### Description Regarding issue 113 The PR adds tests for SFRemoteStorageUtil to increase code coverage ### Checklist - [x] Code compiles correctly - [x] Code is formatted according to [Coding Conventions](../CodingConventions.md) - [x] Created tests which fail without the change (if possible) - [x] All tests passing (`dotnet test`) - [x] Extended the README / documentation, if necessary - [x] Provide JIRA issue id (if possible) or GitHub issue id in PR name --- .../Mock/MockRemoteStorageClient.cs | 148 +++++ .../UnitTests/SFRemoteStorageClientTest.cs | 574 ++++++++++++++++++ .../StorageClient/SFRemoteStorageUtil.cs | 28 +- 3 files changed, 737 insertions(+), 13 deletions(-) create mode 100644 Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs create mode 100644 Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs diff --git a/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs new file mode 100644 index 000000000..c93173d58 --- /dev/null +++ b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs @@ -0,0 +1,148 @@ +/* + * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved. + */ + +using Moq; +using Snowflake.Data.Core.FileTransfer.StorageClient; +using System; +using System.IO; +using System.Net; +using System.Text; + +namespace Snowflake.Data.Tests.Mock +{ + + class MockRemoteStorageClient + { + // Mock data for downloaded file + internal const string FileContent = "RemoteStorageClientTest"; + + // Mock content length + internal const int ContentLength = 9999; + + // Mock error message + internal const string ErrorMessage = "Mock GCS Remote Storage Error"; + + // Variables for the encryption data + [ThreadStatic] static Stream t_encryptedStream = null; + [ThreadStatic] static string t_encryptedStreamIV = null; + [ThreadStatic] static string t_encryptedStreamKey = null; + + static internal void SetEncryptionData(Stream stream, string iv, string key) + { + t_encryptedStream = stream; + t_encryptedStreamIV = iv; + t_encryptedStreamKey = key; + } + + // Sets up the mock sequence when Remote Storage uploads a file + // 1. Get the file header + // 2. Upload the file + static internal HttpWebResponse SequenceResponseForUploadFile(ref bool firstRequest, HttpStatusCode statusCode, HttpStatusCode statusCodeAfterRetry) + { + try + { + // For the first call to GetResponse(), return the file header response + // For the second call to GetResponse(), return the upload response + var response = firstRequest ? + CreateResponseForFileHeader(statusCode) : + CreateResponseForUploadFile(statusCodeAfterRetry); + + return response; + } + finally + { + // If an exception is thrown, the value will still be changed + firstRequest = false; + } + } + + // Create a mock response for GetFileHeader + static internal HttpWebResponse CreateResponseForFileHeader(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode == HttpStatusCode.OK) + { + response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + response.Object.Headers.Add("content-length", MockGCSClient.ContentLength.ToString()); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_SFC_DIGEST, MockGCSClient.SFCDigest); + } + else + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + + // Create a mock response for UploadFile + static internal HttpWebResponse CreateResponseForUploadFile(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode != HttpStatusCode.OK) + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + + // Create a mock response for DownloadFile + static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode == HttpStatusCode.OK) + { + response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + + // For downloads with encryption material + if (t_encryptedStream != null) + { + // Set the position to 0 and return the encrypted stream + t_encryptedStream.Position = 0; + response.Setup(c => c.GetResponseStream()).Returns(t_encryptedStream); + + // Set the iv and key to the ones used for encrypting the stream + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, + "{" + + $"\"ContentEncryptionIV\": \"{t_encryptedStreamIV}\", " + + $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{t_encryptedStreamKey}\"}}" + + "}"); + + t_encryptedStreamIV = null; + t_encryptedStreamKey = null; + t_encryptedStream = null; + } + else // For unencrypted downloads + { + response.Setup(c => c.GetResponseStream()).Returns(new MemoryStream(Encoding.ASCII.GetBytes(FileContent))); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, + "{" + + $"\"ContentEncryptionIV\": \"{MockGCSClient.GcsIV}\", " + + $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{MockGCSClient.GcsKey}\"}}" + + "}"); + } + + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_MATDESC_KEY, MockGCSClient.GcsMatdesc); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_SFC_DIGEST, MockGCSClient.SFCDigest); + response.Object.Headers.Add(SFGCSClient.GCS_FILE_HEADER_CONTENT_LENGTH, MockGCSClient.ContentLength.ToString()); + } + else + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + } +} + diff --git a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs new file mode 100644 index 000000000..fa7df6250 --- /dev/null +++ b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs @@ -0,0 +1,574 @@ +/* + * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved. + */ + +namespace Snowflake.Data.Tests.UnitTests +{ + using NUnit.Framework; + using Snowflake.Data.Core; + using Snowflake.Data.Core.FileTransfer.StorageClient; + using Snowflake.Data.Core.FileTransfer; + using System.Collections.Generic; + using System; + using Snowflake.Data.Tests.Mock; + using System.Threading.Tasks; + using System.Threading; + using System.IO; + using System.Text; + using System.Net; + using Moq; + + [TestFixture] + class SFRemoteStorageClientTest : SFBaseTest + { + // Mock data for file metadata + const string EndPoint = "mockEndPoint.com"; + + const string LocationStage = "mock-customer-stage"; + const string LocationId = "mock-id"; + const string LocationTables = "tables"; + const string LocationKey = "mock-key"; + const string LocationPath = LocationTables + "/" + LocationKey + "/"; + const string Location = LocationStage + "/" + LocationId + "/" + LocationPath; + + const string Region = "us-west-2"; + + const string StorageAccount = "mockStorageAccount"; + + [ThreadStatic] private static string t_realSourceFilePath; + + PutGetEncryptionMaterial EncryptionMaterial = new PutGetEncryptionMaterial() + { + queryId = "MOCK/QUERY/ID/==", + queryStageMasterKey = "MOCKQUERYSTAGEMASTERKE==", + smkId = 9999 + }; + + // Mock unsupported stage type + const string UnsupportedStageType = "UNSUPPORTED"; + + // Settings for mock client + const int Parallel = 0; + + // File name for mock test files + [ThreadStatic] private static string t_downloadFileName; + const string LocalLocation = "./"; + + // Mock upload file size + const int UploadFileSize = 9999; + const int DestFileSizeWhenFileAlreadyExists = 0; + + // Mock client and metadata + SFGCSClient _client; + PutGetResponseData _responseData; + SFFileMetadata _fileMetadata; + + // Token for async tests + CancellationToken _cancellationToken; + + // Flags for non-async and async mock methods + const bool NotAsync = false; + const bool IsAsync = true; + + [SetUp] + public void BeforeTest() + { + t_realSourceFilePath = TestNameWithWorker + "_realSrcFilePath.txt"; + t_downloadFileName = TestNameWithWorker + "_mockFileName.txt"; + + _fileMetadata = new SFFileMetadata() + { + destFileName = t_downloadFileName, + localLocation = LocalLocation, + MaxBytesInMemory = 1024, + memoryStream = new MemoryStream(Encoding.UTF8.GetBytes(MockRemoteStorageClient.FileContent)), + parallel = Parallel, + realSrcFilePath = t_realSourceFilePath, + stageInfo = new PutGetStageInfo() + { + endPoint = EndPoint, + isClientSideEncrypted = true, + location = Location, + locationType = SFRemoteStorageUtil.GCS_FS, + path = LocationPath, + presignedUrl = null, + region = Region, + stageCredentials = new Dictionary() + { + {"AWS_KEY_ID", "MOCK_AWS_KEY_ID"}, + {"AWS_SECRET_KEY", "MOCK_AWS_SECRET_KEY"}, + {"AWS_TOKEN", "MOCK_AWS_TOKEN"}, + {"AWS_ID", "MOCK_AWS_ID"}, + {"AWS_KEY", "MOCK_AWS_KEY"}, + {"AZURE_SAS_TOKEN", "MOCK_AZURE_SAS_TOKEN"}, + {"GCS_ACCESS_TOKEN", "MOCK_GCS_ACCESS_TOKEN"} + }, + storageAccount = StorageAccount + }, + uploadSize = UploadFileSize + }; + + _responseData = new PutGetResponseData() + { + autoCompress = true, + stageInfo = _fileMetadata.stageInfo + }; + + // Set the mock GCS client to use + _client = new SFGCSClient(_fileMetadata.stageInfo); + _fileMetadata.client = _client; + + _cancellationToken = new CancellationToken(); + } + + [TearDown] + public void AfterTest() + { + // Delete temporary files from upload + if (File.Exists(_fileMetadata.realSrcFilePath)) + { + File.Delete(_fileMetadata.realSrcFilePath); + } + + // Delete temporary files from download + if (File.Exists(t_downloadFileName)) + { + File.Delete(t_downloadFileName); + } + } + + [Test] + [Ignore("RemoteStorageClientTest")] + public void RemoteStorageClientTestDone() + { + // Do nothing; + } + + [Test] + [TestCase(SFRemoteStorageUtil.LOCAL_FS)] + [TestCase(SFRemoteStorageUtil.S3_FS)] + [TestCase(SFRemoteStorageUtil.AZURE_FS)] + [TestCase(SFRemoteStorageUtil.GCS_FS)] + [TestCase(UnsupportedStageType)] // Any other stage type should return null + public void TestGetRemoteStorageClient(string stageType) + { + _responseData.stageInfo.locationType = stageType; + + if (stageType == SFRemoteStorageUtil.LOCAL_FS) + { + Assert.Throws(() => SFRemoteStorageUtil.GetRemoteStorage(_responseData)); + } + else + { + ISFRemoteStorageClient client = SFRemoteStorageUtil.GetRemoteStorage(_responseData); + + if (stageType == SFRemoteStorageUtil.S3_FS) + { + Assert.IsInstanceOf(client); + } + else if (stageType == SFRemoteStorageUtil.AZURE_FS) + { + Assert.IsInstanceOf(client); + } + else if (stageType == SFRemoteStorageUtil.GCS_FS) + { + Assert.IsInstanceOf(client); + } + else + { + Assert.IsNull(client); + } + } + } + + [Test] + [TestCase(false, false)] + [TestCase(false, true)] + [TestCase(true, false)] + [TestCase(true, true)] + public void TestUploadFileOrStreamWithAndWithoutEncryption(bool containsEncryptionMaterial, bool useMemoryStream) + { + // Arrange + var mockWebRequest = new Mock(); + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.CreateResponseForFileHeader(HttpStatusCode.OK); + }); + _client.SetCustomWebRequest(mockWebRequest.Object); + + // Add encryption material to the file metadata + _fileMetadata.encryptionMaterial = containsEncryptionMaterial ? EncryptionMaterial : null; + + // Use file for upload + if (!useMemoryStream) + { + File.WriteAllText(_fileMetadata.realSrcFilePath, MockRemoteStorageClient.FileContent); + _fileMetadata.memoryStream = null; + } + + // Act + SFRemoteStorageUtil.UploadOneFile(_fileMetadata); + + // Assert + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + Assert.AreEqual(ResultStatus.SKIPPED.ToString(), _fileMetadata.resultStatus); + } + + private Mock CreateBaseMockClient() + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = new Mock(); + mockWebRequest.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + + return mockWebRequest; + } + + private void SetUpMockClientForUpload(HttpStatusCode statusCode, HttpStatusCode statusCodeAfterRetry, bool isAsync) + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = CreateBaseMockClient(); + bool firstRequest = true; + + if (isAsync) + { + mockWebRequest.Setup(client => client.GetRequestStreamAsync()) + .Returns(() => Task.FromResult(new MemoryStream())); + mockWebRequest.Setup(client => client.GetResponseAsync()) + .Returns(() => + { + return Task.FromResult(MockRemoteStorageClient.SequenceResponseForUploadFile(ref firstRequest, statusCode, statusCodeAfterRetry)); + }); + } + else + { + mockWebRequest.Setup(client => client.GetRequestStream()) + .Returns(() => new MemoryStream()); + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.SequenceResponseForUploadFile(ref firstRequest, statusCode, statusCodeAfterRetry); + }); + } + + _client.SetCustomWebRequest(mockWebRequest.Object); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + public void TestUploadOneFileWithRetry(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + SFRemoteStorageUtil.UploadOneFileWithRetry(_fileMetadata); + + // Assert + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + public async Task TestUploadOneFileAsyncWithRetry(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + await SFRemoteStorageUtil.UploadOneFileWithRetryAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, null, ResultStatus.SKIPPED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.BadRequest, ResultStatus.RENEW_PRESIGNED_URL)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public void TestUploadOneFile(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + SFRemoteStorageUtil.UploadOneFile(_fileMetadata); + + // Assert + if (expectedResultStatus == ResultStatus.SKIPPED) + { + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, null, ResultStatus.SKIPPED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.BadRequest, ResultStatus.RENEW_PRESIGNED_URL)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public async Task TestUploadOneFileAsync(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + if (expectedResultStatus == ResultStatus.SKIPPED) + { + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestUploadOneFileThrowsForRetryErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.UploadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestUploadOneFileAsyncThrowsForRetryErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + + [Test] + [TestCase(HttpStatusCode.NotFound, null, ResultStatus.ERROR)] + public void TestUploadOneFileThrowsForUnknownErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.UploadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in uploading a file: .*")); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, null, ResultStatus.ERROR)] + public void TestUploadOneFileAsyncThrowsForUnknownErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in uploading a file: .*")); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + private void SetUpMockClientForDownload(HttpStatusCode statusCode, bool isAsync) + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = CreateBaseMockClient(); + if (isAsync) + { + mockWebRequest.Setup(client => client.GetResponseAsync()) + .Returns(() => + { + return Task.FromResult(MockRemoteStorageClient.CreateResponseForDownloadFile(statusCode)); + }); + + } + else + { + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.CreateResponseForDownloadFile(statusCode); + }); + } + + _client.SetCustomWebRequest(mockWebRequest.Object); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + [TestCase(HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public void TestDownloadOneFile(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + SFRemoteStorageUtil.DownloadOneFile(_fileMetadata); + + // Assert + if (expectedResultStatus == ResultStatus.DOWNLOADED) + { + string text = File.ReadAllText(t_downloadFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + [TestCase(HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public async Task TestDownloadOneFileAsync(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + if (expectedResultStatus == ResultStatus.DOWNLOADED) + { + string text = File.ReadAllText(t_downloadFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestDownloadOneFileThrowsForRetryErrors(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.DownloadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestDownloadOneFileAsyncThrowsForRetryErrors(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound)] + public void TestDownloadOneFileThrowsForUnknownErrors(HttpStatusCode httpStatusCode) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.DownloadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in downloading a file: .*")); + Assert.IsNull(_fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound)] + public void TestDownloadOneFileAsyncThrowsForUnknownErrors(HttpStatusCode httpStatusCode) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in downloading a file: .*")); + Assert.IsNull(_fileMetadata.resultStatus); + } + + private void SetUpMockEncryptedFileForDownload() + { + _fileMetadata.encryptionMaterial = EncryptionMaterial; + + // Write file to encrypt + File.WriteAllText(t_downloadFileName, MockRemoteStorageClient.FileContent); + + // Get encrypted stream from file + SFEncryptionMetadata encryptionMetadata = new SFEncryptionMetadata(); + Stream stream = EncryptionProvider.EncryptFile( + t_downloadFileName, + _fileMetadata.encryptionMaterial, + encryptionMetadata, + FileTransferConfiguration.FromFileMetadata(_fileMetadata)); + + // Set up the stream and metadata for decryption + MockRemoteStorageClient.SetEncryptionData(stream, encryptionMetadata.iv, encryptionMetadata.key); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + public void TestDownloadOneFileWithEncryptionMaterial(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockEncryptedFileForDownload(); + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + SFRemoteStorageUtil.DownloadOneFile(_fileMetadata); + + // Assert + string text = File.ReadAllText(t_downloadFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + public async Task TestDownloadOneFileAsyncWithEncryptionMaterial(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockEncryptedFileForDownload(); + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + string text = File.ReadAllText(t_downloadFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + } +} diff --git a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs index 45489954f..92e8dd467 100644 --- a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs +++ b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs @@ -377,6 +377,11 @@ internal static void DownloadOneFile(SFFileMetadata fileMetadata) fileMetadata.destFileSize = fileInfo.Length; return; } + else if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || + fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) + { + return; + } else { HandleDownloadFileErr(ref fileMetadata, ref maxConcurrency, ref lastErr, retry, maxRetry); @@ -410,12 +415,6 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can } ISFRemoteStorageClient client = fileMetadata.client; - FileHeader fileHeader = await client.GetFileHeaderAsync(fileMetadata, cancellationToken).ConfigureAwait(false); - - if (fileHeader != null) - { - fileMetadata.srcFileSize = fileHeader.contentLength; - } int maxConcurrency = fileMetadata.parallel; Exception lastErr = null; @@ -440,15 +439,18 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can * One example of this is the utils that use presigned url * for upload / download and not the storage client library. **/ + FileHeader fileHeader = null; if (fileMetadata.presignedUrl != null) { fileHeader = await client.GetFileHeaderAsync(fileMetadata, cancellationToken).ConfigureAwait(false); } + SFEncryptionMetadata encryptionMetadata = fileHeader != null ? fileHeader.encryptionMetadata : fileMetadata.encryptionMetadata; + string tmpDstName = EncryptionProvider.DecryptFile( fullDstPath, fileMetadata.encryptionMaterial, - fileHeader.encryptionMetadata, + encryptionMetadata, FileTransferConfiguration.FromFileMetadata(fileMetadata)); File.Delete(fullDstPath); @@ -464,6 +466,11 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can fileMetadata.destFileSize = fileInfo.Length; return; } + else if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || + fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) + { + return; + } else { HandleDownloadFileErr(ref fileMetadata, ref maxConcurrency, ref lastErr, retry, maxRetry); @@ -490,12 +497,7 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can /// The max retry private static void HandleDownloadFileErr(ref SFFileMetadata fileMetadata, ref int maxConcurrency, ref Exception lastErr, int retry, int maxRetry) { - if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || - fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) - { - return; - } - else if (fileMetadata.resultStatus == ResultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY.ToString()) + if (fileMetadata.resultStatus == ResultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY.ToString()) { lastErr = fileMetadata.lastError; // Failed to download file, retrying with max concurrency