From 9abe323019c5d762490ed8564420ccc2e5561372 Mon Sep 17 00:00:00 2001
From: Tony Li <tony.li@automattic.com>
Date: Mon, 4 Mar 2024 13:32:17 +1300
Subject: [PATCH] Add an unit test where HTTP response contains a large amount
 of data

---
 .../Utilities/URLSessionHelperTests.swift     | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/WordPressKitTests/Utilities/URLSessionHelperTests.swift b/WordPressKitTests/Utilities/URLSessionHelperTests.swift
index 999af86f..11371a26 100644
--- a/WordPressKitTests/Utilities/URLSessionHelperTests.swift
+++ b/WordPressKitTests/Utilities/URLSessionHelperTests.swift
@@ -1,4 +1,5 @@
 import Foundation
+import CryptoKit
 import XCTest
 import OHHTTPStubs
 
@@ -235,6 +236,54 @@ class URLSessionHelperTests: XCTestCase {
         let expectedBody = "--\(boundary)\r\nContent-Disposition: form-data; name=\"name\"\r\n\r\nvalue\r\n--\(boundary)--\r\n"
         XCTAssertEqual(String(data: requestBody, encoding: .utf8), expectedBody)
     }
+
+    func testGetLargeData() async throws {
+        let file = try self.createLargeFile(megaBytes: 100)
+        defer {
+            try? FileManager.default.removeItem(at: file)
+        }
+
+        stub(condition: isPath("/hello")) { _ in
+            HTTPStubsResponse(fileURL: file, statusCode: 200, headers: nil)
+        }
+
+        let builder = HTTPRequestBuilder(url: URL(string: "https://wordpress.org/hello")!)
+        let response = try await session.perform(request: builder, errorType: TestError.self).get()
+
+        try XCTAssertEqual(
+            sha256(XCTUnwrap(InputStream(url: file))),
+            sha256(InputStream(data: response.body))
+        )
+    }
+
+    private func createLargeFile(megaBytes: Int) throws -> URL {
+        let fileManager = FileManager.default
+        let file = try fileManager.url(for: .documentDirectory, in: .userDomainMask, appropriateFor: nil, create: true)
+            .appendingPathComponent("large-file-\(UUID().uuidString).txt")
+        fileManager.createFile(atPath: file.path, contents: nil)
+
+        let handle = try FileHandle(forUpdating: file)
+        for _ in 0..<megaBytes {
+            handle.write(Data(repeating: 46, count: 1_000_000))
+        }
+        try handle.close()
+        return file
+    }
+
+    private func sha256(_ stream: InputStream) -> SHA256Digest {
+        stream.open()
+        defer { stream.close() }
+
+        var hash = SHA256()
+        let maxLength = 1024
+        var buffer = [UInt8](repeating: 0, count: maxLength)
+        while stream.hasBytesAvailable {
+            let bytes = stream.read(&buffer, maxLength: maxLength)
+            let data = Data(bytesNoCopy: &buffer, count: bytes, deallocator: .none)
+            hash.update(data: data)
+        }
+        return hash.finalize()
+    }
 }
 
 class BackgroundURLSessionHelperTests: URLSessionHelperTests {