More flexible checksum test refactors

awslabs · Mar 12, 2024 · 7dd0ff9 · 7dd0ff9
1 parent fc51f47
commit 7dd0ff9
Showing 1 changed file with 59 additions and 60 deletions.
diff --git a/IntegrationTests/Services/AWSS3IntegrationTests/S3FlexibleChecksumsTests.swift b/IntegrationTests/Services/AWSS3IntegrationTests/S3FlexibleChecksumsTests.swift
@@ -17,57 +17,9 @@ final class S3FlexibleChecksumsTests: S3XCTestCase {
         super.setUp()
         oneMBData = Data((0..<(1024 * 1024)).map { _ in UInt8.random(in: UInt8.min...UInt8.max) })
     }
+
+    // MARK: - Data uploads
 
-    private func _testPutGetObject(withChecksumAlgorithm algorithm: S3ClientTypes.ChecksumAlgorithm, objectNameSuffix: String, upload: ByteStream, file: StaticString = #filePath, line: UInt = #line) async throws {
-        let objectName = "flexible-checksums-s3-test-\(objectNameSuffix)"
-
-        let input = PutObjectInput(
-            body: upload,
-            bucket: bucketName,
-            checksumAlgorithm: algorithm,
-            key: objectName
-        )
-
-        let output = try await client.putObject(input: input)
-
-        // Verify the checksum response based on the algorithm used.
-        let checksumResponse = try XCTUnwrap(getChecksumResponse(from: output, with: algorithm), file: file, line: line)
-        XCTAssertNotNil(checksumResponse, file: file, line: line)
-
-        let getInput = GetObjectInput(bucket: bucketName, checksumMode: S3ClientTypes.ChecksumMode.enabled, key: objectName)
-        let getOutput = try await client.getObject(input: getInput) // will error for normal payloads if checksum mismatch
-        XCTAssertNotNil(getOutput.body, file: file, line: line) // Ensure there's a body in the response.
-
-        // Additional step for stream: Validate stream and read data.
-        if case .stream = upload {
-            let streamingBody = try XCTUnwrap(getOutput.body, file: file, line: line)
-            if case .stream(let stream) = streamingBody {
-                XCTAssert(stream is ValidatingBufferedStream, "Expected ValidatingBufferedStream for streaming upload", file: file, line: line)
-                let data = try await streamingBody.readData() // will error if checksum mismatch
-                XCTAssertNotNil(data, file: file, line: line)
-            } else {
-                XCTFail("Did not receive a stream when expected for checksum validation!", file: file, line: line)
-            }
-        }
-    }
-
-    private func getChecksumResponse(from response: PutObjectOutput, with algorithm: S3ClientTypes.ChecksumAlgorithm) throws -> String? {
-        switch algorithm {
-        case .crc32:
-            return response.checksumCRC32
-        case .crc32c:
-            return response.checksumCRC32C
-        case .sha1:
-            return response.checksumSHA1
-        case .sha256:
-            return response.checksumSHA256
-        default:
-            XCTFail("Unsupported checksum algorithm")
-            return nil
-        }
-    }
-
-    // Test cases for data uploads
     func test_putGetObject_data_crc32() async throws {
         try await _testPutGetObject(withChecksumAlgorithm: .crc32, objectNameSuffix: "crc32-data", upload: .data(oneMBData))
     }
@@ -84,7 +36,8 @@ final class S3FlexibleChecksumsTests: S3XCTestCase {
         try await _testPutGetObject(withChecksumAlgorithm: .sha256, objectNameSuffix: "sha256-data", upload: .data(oneMBData))
     }
 
-    // Test cases for streaming uploads
+    // MARK: - Streaming uploads
+
     func test_putGetObject_streaming_crc32() async throws {
         let bufferedStream = BufferedStream(data: oneMBData, isClosed: true)
         try await _testPutGetObject(withChecksumAlgorithm: .crc32, objectNameSuffix: "crc32", upload: .stream(bufferedStream))
@@ -110,21 +63,67 @@ final class S3FlexibleChecksumsTests: S3XCTestCase {
         let bufferedStream = BufferedStream(data: oneMBData, isClosed: true)
         let objectName = "flexible-checksums-s3-test-chunked"
 
-        let input = PutObjectInput(
+        let putObjectInput = PutObjectInput(
             body: .stream(bufferedStream),
             bucket: bucketName,
             key: objectName
         )
 
-        do {
-            let response = try await client.putObject(input: input)
-            try XCTAssertNotNil(XCTUnwrap(response))
+        _ = try await client.putObject(input: putObjectInput)
+
+        let getObjectInput = GetObjectInput(bucket: bucketName, key: objectName)
+        let getObjectOutput = try await client.getObject(input: getObjectInput)
+        XCTAssertNotNil(getObjectOutput.body) // Ensure there's a body in the response.
+    }
+
+    // MARK: - Private methods
+
+    private func _testPutGetObject(withChecksumAlgorithm algorithm: S3ClientTypes.ChecksumAlgorithm, objectNameSuffix: String, upload: ByteStream, file: StaticString = #filePath, line: UInt = #line) async throws {
+        let objectName = "flexible-checksums-s3-test-\(objectNameSuffix)"
+
+        let input = PutObjectInput(
+            body: upload,
+            bucket: bucketName,
+            checksumAlgorithm: algorithm,
+            key: objectName
+        )
+
+        let output = try await client.putObject(input: input)
+
+        // Verify the checksum response based on the algorithm used.
+        let checksumResponse = try XCTUnwrap(getChecksumResponse(from: output, with: algorithm), file: file, line: line)
+        XCTAssertNotNil(checksumResponse, file: file, line: line)
+
+        let getInput = GetObjectInput(bucket: bucketName, checksumMode: S3ClientTypes.ChecksumMode.enabled, key: objectName)
+        let getOutput = try await client.getObject(input: getInput) // will error for normal payloads if checksum mismatch
+        XCTAssertNotNil(getOutput.body, file: file, line: line) // Ensure there's a body in the response.
+
+        // Additional step for stream: Validate stream and read data.
+        if case .stream = upload {
+            let streamingBody = try XCTUnwrap(getOutput.body, file: file, line: line)
+            if case .stream(let stream) = streamingBody {
+                XCTAssert(stream is ValidatingBufferedStream, "Expected ValidatingBufferedStream for streaming upload", file: file, line: line)
+                let data = try await streamingBody.readData() // will error if checksum mismatch
+                XCTAssertNotNil(data, file: file, line: line)
+            } else {
+                XCTFail("Did not receive a stream when expected for checksum validation!", file: file, line: line)
+            }
+        }
+    }
 
-            let getInput = GetObjectInput(bucket: bucketName, key: objectName)
-            let responseGet = try await client.getObject(input: getInput)
-            XCTAssertNotNil(responseGet.body) // Ensure there's a body in the response.
-        } catch {
-            throw error
+    private func getChecksumResponse(from response: PutObjectOutput, with algorithm: S3ClientTypes.ChecksumAlgorithm) throws -> String? {
+        switch algorithm {
+        case .crc32:
+            return response.checksumCRC32
+        case .crc32c:
+            return response.checksumCRC32C
+        case .sha1:
+            return response.checksumSHA1
+        case .sha256:
+            return response.checksumSHA256
+        default:
+            XCTFail("Unsupported checksum algorithm")
+            return nil
         }
     }
 }