LocalData/Tests/LocalDataTests/EncryptionHelperTests.swift

import Foundation
import Testing
@testable import LocalData

struct EncryptionHelperTests {
    private let masterKeyService = "LocalData"
    private let keyName = "LocalDataTests.encryption"
    private let payload = Data("payload".utf8)
    
    @Test func aesGCMWithPBKDF2RoundTrip() async throws {
        await clearMasterKey()
        
        let policy: SecurityPolicy.EncryptionPolicy = .aes256(
            keyDerivation: .pbkdf2(iterations: 1_000)
        )
        
        let encrypted = try await EncryptionHelper.shared.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await EncryptionHelper.shared.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
        await clearMasterKey()
    }
    
    @Test func chaChaPolyWithHKDFRoundTrip() async throws {
        await clearMasterKey()
        
        let policy: SecurityPolicy.EncryptionPolicy = .chacha20Poly1305(
            keyDerivation: .hkdf()
        )
        
        let encrypted = try await EncryptionHelper.shared.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await EncryptionHelper.shared.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
        await clearMasterKey()
    }

    @Test func customConfigurationRoundTrip() async throws {
        let customService = "Test.CustomService"
        let customAccount = "Test.CustomAccount"
        let config = EncryptionConfiguration(
            masterKeyService: customService,
            masterKeyAccount: customAccount
        )
        
        await EncryptionHelper.shared.updateConfiguration(config)
        
        let policy: SecurityPolicy.EncryptionPolicy = .chacha20Poly1305(
            keyDerivation: .hkdf()
        )
        
        let encrypted = try await EncryptionHelper.shared.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await EncryptionHelper.shared.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
        
        // Cleanup keychain
        try? await KeychainHelper.shared.deleteAll(service: customService)
        
        // Reset to default
        await EncryptionHelper.shared.updateConfiguration(.default)
    }
    
    @Test func externalProviderWithHKDFRoundTrip() async throws {
        let source = KeyMaterialSource(id: "test.external")
        let provider = StaticKeyMaterialProvider(material: Data(repeating: 7, count: 32))
        await EncryptionHelper.shared.registerKeyMaterialProvider(provider, for: source)

        let policy: SecurityPolicy.EncryptionPolicy = .external(
            source: source,
            keyDerivation: .hkdf()
        )

        let encrypted = try await EncryptionHelper.shared.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await EncryptionHelper.shared.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )

        #expect(decrypted == payload)
    }
    
    private func clearMasterKey() async {
        try? await KeychainHelper.shared.deleteAll(service: masterKeyService)
    }
}

private struct StaticKeyMaterialProvider: KeyMaterialProviding {
    let material: Data

    func keyMaterial(for keyName: String) async throws -> Data {
        material
    }
}