passforios/passKitTests/LowLevel/PGPAgentLowLevelTests.swift

//
//  PGPAgentLowLevelTests.swift
//  passKitTests
//
//  Detailed unit tests tracking the exact API call behavior of PGPAgent.decrypt.
//  Uses MockPGPInterface to verify what arguments are passed to the underlying
//  PGPInterface methods, and how passphrase resolution interacts with the keystore
//  and the requestPGPKeyPassphrase callback.
//

import XCTest

@testable import passKit

final class PGPAgentLowLevelTests: XCTestCase {
    private var keychain: DictBasedKeychain!
    private var mockPGP: MockPGPInterface!
    private var agent: PGPAgent!

    private let testEncryptedData = Data("encrypted-payload".utf8)
    private let testDecryptedData = Data("decrypted-payload".utf8)

    /// Tracks all calls to requestPGPKeyPassphrase closures created via `passphraseCallback(_:)`.
    private var passphraseRequests: [String] = []

    /// Creates a requestPGPKeyPassphrase closure that records the keyID it's called with
    /// into `passphraseRequests` and returns `response`.
    private func passphraseCallback(_ response: String) -> (String) -> String {
        { [self] keyID in
            passphraseRequests.append(keyID)
            return response
        }
    }

    override func setUp() {
        super.setUp()

        keychain = DictBasedKeychain()
        // Set pgpKeyPassphrase key so checkAndInit() doesn't re-init and overwrite our mock.
        keychain.add(string: "dummy", for: Globals.pgpKeyPassphrase)

        mockPGP = MockPGPInterface()
        // some defaults
        mockPGP.decryptResult = testDecryptedData
        mockPGP.encryptResult = Data("mock-encrypted".utf8)

        passphraseRequests = []

        agent = PGPAgent(keyStore: keychain, pgpInterface: mockPGP)
    }

    override func tearDown() {
        keychain.removeAllContent()
        super.tearDown()
    }

    // MARK: - decrypt(encryptedData:keyID:requestPGPKeyPassphrase:) - Key Resolution

    /// When the private key is found, decrypt is called with the provided keyID.
    func testDecryptWithKeyID_keyFound_usesProvidedKeyID() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        let result = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(result, testDecryptedData)
        XCTAssertEqual(mockPGP.decryptCalls.count, 1)
        XCTAssertEqual(mockPGP.decryptCalls[0].keyID, longFingerprint)
        XCTAssertEqual(mockPGP.decryptCalls[0].encryptedData, testEncryptedData)
    }

    /// When the private key is NOT found but there's exactly one key, falls back to that key.
    func testDecryptWithKeyID_keyNotFound_singleKey_fallsBackToOnlyKey() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [] // requested key not found
        mockPGP.keyIDs = [longFingerprint]

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: "UNKNOWN", requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.decryptCalls.count, 1)
        // The keyID passed to pgpInterface.decrypt should be the fallback key, not the requested one.
        XCTAssertEqual(mockPGP.decryptCalls[0].keyID, longFingerprint)
    }

    /// When the private key is NOT found and there are multiple keys, throws pgpPrivateKeyNotFound.
    func testDecryptWithKeyID_keyNotFound_multipleKeys_throws() {
        mockPGP.privateKeyIDs = []
        mockPGP.keyIDs = ["4712286271220db299883ea7062e678da1024dae", "787eae1a5fa3e749aa34cc6aa0645ebed862027e"]

        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, keyID: "UNKNOWN", requestPGPKeyPassphrase: passphraseCallback("pass"))) { error in
            XCTAssertEqual(error as? AppError, AppError.pgpPrivateKeyNotFound(keyID: "UNKNOWN"))
        }
        // pgpInterface.decrypt should NOT have been called
        XCTAssertEqual(mockPGP.decryptCalls.count, 0)
    }

    /// containsPrivateKey is called with the provided keyID to check membership.
    func testDecryptWithKeyID_checksContainsPrivateKey() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.containsPrivateKeyCalls, [shortID])
    }

    // MARK: - decrypt(encryptedData:keyID:requestPGPKeyPassphrase:) - Passphrase Resolution

    /// On first decrypt (latestDecryptStatus=true), the passphrase is looked up from keystore first.
    /// If found in keystore, requestPGPKeyPassphrase is NOT called.
    func testDecryptWithKeyID_firstCall_passphraseFromKeystore() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.selectedKeyForPassphrase = longFingerprint
        keychain.add(string: "stored-passphrase", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("requested-passphrase"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["stored-passphrase"])
        XCTAssertEqual(passphraseRequests, [], "requestPGPKeyPassphrase should not be called when passphrase is in keystore")
    }

    /// On first decrypt, if keystore doesn't have the passphrase, requestPGPKeyPassphrase is called.
    /// The keyID passed to requestPGPKeyPassphrase is the (possibly resolved) keyID.
    func testDecryptWithKeyID_firstCall_passphraseFromRequest() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.selectedKeyForPassphrase = shortID
        // No passphrase in keystore for this key.
        XCTAssertFalse(keychain.contains(key: AppKeychain.getPGPKeyPassphraseKey(keyID: shortID)))
        XCTAssertFalse(keychain.contains(key: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint)))

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("my-passphrase"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["my-passphrase"])
        XCTAssertEqual(passphraseRequests, [shortID])
    }

    /// After a failed decrypt (latestDecryptStatus=false), requestPGPKeyPassphrase is ALWAYS called,
    /// even if the keystore has a cached passphrase.
    func testDecryptWithKeyID_afterFailure_alwaysRequestsPassphrase() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        keychain.add(string: "stored-passphrase", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))

        // First call: force a failure by making decrypt throw.
        mockPGP.decryptError = AppError.wrongPassphrase
        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("bad")))

        // Now latestDecryptStatus=false. Second call should always request.
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = longFingerprint
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("fresh-passphrase"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh-passphrase"])
        XCTAssertEqual(passphraseRequests, [longFingerprint], "After failure, passphrase should always be requested")
    }

    /// After a successful decrypt, the next call uses keystore first (latestDecryptStatus=true).
    func testDecryptWithKeyID_afterSuccess_usesKeystoreFirst() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        // First call succeeds.
        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass1"))

        // Store a passphrase in keystore under the short ID (matching what PGPAgent used for lookup).
        keychain.add(string: "pass1", for: AppKeychain.getPGPKeyPassphraseKey(keyID: shortID))

        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = shortID
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("ignored-passphrase"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["pass1"])
        XCTAssertEqual(passphraseRequests, [])
    }

    // MARK: - decrypt(encryptedData:keyID:requestPGPKeyPassphrase:) - Return Values & Error Propagation

    /// When pgpInterface.decrypt returns nil, agent.decrypt returns nil.
    func testDecryptWithKeyID_interfaceReturnsNil_returnsNil() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.decryptResult = nil

        let result = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertNil(result)
    }

    /// When pgpInterface.decrypt returns nil, latestDecryptStatus stays false
    /// (next call will always request passphrase).
    func testDecryptWithKeyID_interfaceReturnsNil_statusStaysFalse() throws {
        let shortID = "d862027e"
        let longFingerprint = "787eae1a5fa3e749aa34cc6aa0645ebed862027e"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.decryptResult = nil

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass"))

        // Second call - should always request (latestDecryptStatus=false because nil return doesn't set it to true).
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: shortID))
        keychain.add(string: "cached-long", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))
        mockPGP.decryptResult = testDecryptedData
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = longFingerprint
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh"])
        XCTAssertEqual(passphraseRequests, [longFingerprint], "After nil return, passphrase should always be requested")
    }

    /// When pgpInterface.decrypt throws, the error propagates and latestDecryptStatus stays false.
    func testDecryptWithKeyID_interfaceThrows_propagatesError() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.decryptError = AppError.wrongPassphrase
        mockPGP.selectedKeyForPassphrase = longFingerprint

        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass"))) { error in
            XCTAssertEqual(error as? AppError, AppError.wrongPassphrase)
        }
        XCTAssertEqual(passphraseRequests, [longFingerprint])

        // Verify latestDecryptStatus stayed false: next call should always request passphrase,
        // even though the keystore has one cached.
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: shortID))
        keychain.add(string: "cached-long", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = longFingerprint
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh"])
        XCTAssertEqual(passphraseRequests, [longFingerprint], "After throw, passphrase should always be requested (latestDecryptStatus=false)")
    }

    /// After successful decrypt, latestDecryptStatus is true.
    func testDecryptWithKeyID_success_setsStatusTrue() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        // Force latestDecryptStatus to false first.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("bad"))
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        passphraseRequests.removeAll()

        // Now succeed.
        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("good"))

        // Third call: latestDecryptStatus=true, so should try keystore first.
        keychain.add(string: "good", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = longFingerprint
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("should-not-use"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["good"])
        XCTAssertEqual(passphraseRequests, [], "After success, should try keystore first")
    }

    // MARK: - decrypt(encryptedData:keyID:requestPGPKeyPassphrase:) - checkAndInit behavior

    /// checkAndInit re-initializes if pgpKeyPassphrase is missing from keystore.
    /// Since we're using a mock as pgpInterface, initKeys would overwrite it; verify the precondition holds.
    func testDecryptWithKeyID_checkAndInit_requiresPGPKeyPassphraseInKeystore() throws {
        // Remove the pgpKeyPassphrase sentinel, which will trigger checkAndInit -> initKeys.
        keychain.removeContent(for: Globals.pgpKeyPassphrase)
        // initKeys needs real PGP keys, which we don't have. It should throw keyImport.
        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, keyID: "a1024dae", requestPGPKeyPassphrase: passphraseCallback("pass"))) { error in
            XCTAssertEqual(error as? AppError, AppError.keyImport)
        }
        XCTAssertEqual(passphraseRequests, [], "requestPGPKeyPassphrase should not be called when checkAndInit fails")
    }

    // MARK: - decrypt(encryptedData:requestPGPKeyPassphrase:) - No KeyID Overload

    /// The no-keyID overload passes nil as keyID to pgpInterface.decrypt
    func testDecryptNoKeyID_passesNilKeyIDToInterface() throws {
        let result = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(result, testDecryptedData)
        XCTAssertEqual(mockPGP.decryptCalls.count, 1)
        XCTAssertNil(mockPGP.decryptCalls[0].keyID)
    }

    /// After failure, the no-keyID overload always requests passphrase.
    func testDecryptNoKeyID_afterFailure_alwaysRequestsPassphrase() throws {
        // Force a failure.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("bad"))

        // Next one can succeed
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: ""))
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = ""
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh"])
        XCTAssertEqual(passphraseRequests, [""])
    }

    /// The no-keyID overload returns nil when interface returns nil.
    func testDecryptNoKeyID_interfaceReturnsNil_returnsNil() throws {
        mockPGP.decryptResult = nil

        let result = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertNil(result)
    }

    /// The no-keyID overload propagates errors from the interface.
    func testDecryptNoKeyID_interfaceThrows_propagatesError() {
        mockPGP.decryptError = AppError.wrongPassphrase

        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("pass"))) { error in
            XCTAssertEqual(error as? AppError, AppError.wrongPassphrase)
        }
    }

    /// The no-keyID overload sets latestDecryptStatus to true on success,
    /// which affects subsequent passphrase lookups.
    func testDecryptNoKeyID_success_setsStatusTrue() throws {
        // Force failure first.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("bad"))
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        passphraseRequests.removeAll()

        // Succeed.
        _ = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("good"))

        // Third call: should try keystore.
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: ""))
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = ""
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("nope"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["cached"])
        XCTAssertEqual(passphraseRequests, [])
    }

    /// The no-keyID overload doesn't check containsPrivateKey.
    func testDecryptNoKeyID_doesNotCheckPrivateKey() throws {
        _ = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.containsPrivateKeyCalls.count, 0)
    }

    // MARK: - Key resolution error vs decrypt status ordering

    /// When pgpPrivateKeyNotFound is thrown (key not found, multiple keys),
    /// latestDecryptStatus is NOT changed because the error occurs BEFORE the status update.
    func testDecryptWithKeyID_keyNotFound_multipleKeys_doesNotChangeDecryptStatus() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = []
        mockPGP.keyIDs = [longFingerprint, "787eae1a5fa3e749aa34cc6aa0645ebed862027e"]

        // This throws pgpPrivateKeyNotFound without changing latestDecryptStatus.
        XCTAssertThrowsError(try agent.decrypt(encryptedData: testEncryptedData, keyID: "UNKNOWN", requestPGPKeyPassphrase: passphraseCallback("pass")))

        // latestDecryptStatus should still be true (initial value).
        // Next call should try keystore first.
        mockPGP.privateKeyIDs = [longFingerprint]
        keychain.add(string: "cached-pass", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))
        mockPGP.selectedKeyForPassphrase = longFingerprint
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(passphraseRequests, [], "After pgpPrivateKeyNotFound, latestDecryptStatus should be unchanged (still true)")
        XCTAssertEqual(mockPGP.resolvedPassphrases, ["cached-pass"])
    }

    /// After failure + key fallback: passphrase is always requested using the RESOLVED (fallback) keyID.
    func testDecryptWithKeyID_afterFailure_keyFallback_requestsWithResolvedKeyID() throws {
        let shortID = "a1024dae"
        let longFingerprint1 = "4712286271220db299883ea7062e678da1024dae"
        let longFingerprint2 = "5fccb081ab8af48972999e2ae750acbfe9444483"
        mockPGP.privateKeyIDs = [longFingerprint1]

        // Force a failure using a short ID that suffix-matches longFingerprint1.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("bad"))

        // Now try with an unknown key that falls back to a different long fingerprint.
        mockPGP.decryptError = nil
        mockPGP.privateKeyIDs = []
        mockPGP.keyIDs = [longFingerprint2]
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = longFingerprint2
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: "e9444483", requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.decryptCalls[0].keyID, longFingerprint2)
        XCTAssertEqual(mockPGP.resolvedPassphrases, ["pass"])
        XCTAssertEqual(passphraseRequests, [longFingerprint2])
    }

    // MARK: - Cross-overload latestDecryptStatus interaction

    /// latestDecryptStatus is shared between both overloads.
    /// A failure in the keyID overload affects the no-keyID overload.
    func testDecryptStatusSharedBetweenOverloads_failureInKeyIDOverload() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        // Fail via keyID overload.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, keyID: longFingerprint, requestPGPKeyPassphrase: passphraseCallback("bad"))

        // Next call via no-keyID overload should always request.
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: ""))
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = ""
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh"])
        XCTAssertEqual(passphraseRequests, [""], "Failure in keyID overload should affect no-keyID overload")
    }

    /// A failure in the no-keyID overload affects the keyID overload.
    func testDecryptStatusSharedBetweenOverloads_failureInNoKeyIDOverload() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]

        // Fail via no-keyID overload.
        mockPGP.decryptError = AppError.wrongPassphrase
        _ = try? agent.decrypt(encryptedData: testEncryptedData, requestPGPKeyPassphrase: passphraseCallback("bad"))

        // Next call via keyID overload should always request.
        mockPGP.decryptError = nil
        mockPGP.decryptCalls.removeAll()
        mockPGP.resolvedPassphrases.removeAll()
        mockPGP.selectedKeyForPassphrase = shortID
        keychain.add(string: "cached", for: AppKeychain.getPGPKeyPassphraseKey(keyID: shortID))
        passphraseRequests.removeAll()

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("fresh"))

        XCTAssertEqual(mockPGP.resolvedPassphrases, ["fresh"])
        XCTAssertEqual(passphraseRequests, [shortID], "Failure in no-keyID overload should affect keyID overload")
    }

    // MARK: - Short vs long key ID behavior

    /// When caller passes a short ID and containsPrivateKey matches it (via suffix), the short ID
    /// is forwarded to pgpInterface.decrypt.
    func testDecryptWithKeyID_shortIDRecognized_shortIDFlowsThrough() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.keyIDs = [longFingerprint]

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.containsPrivateKeyCalls, [shortID])
        XCTAssertEqual(mockPGP.decryptCalls[0].keyID, shortID)
    }

    /// When caller passes a short ID and containsPrivateKey does NOT match, but there's one key,
    /// the long fingerprint from keyID[0] is forwarded instead.
    func testDecryptWithKeyID_shortIDNotRecognized_singleKey_resolvesToLongFingerprint() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [] // short ID doesn't match
        mockPGP.keyIDs = [longFingerprint]

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("pass"))

        XCTAssertEqual(mockPGP.containsPrivateKeyCalls, [shortID])
        XCTAssertEqual(mockPGP.decryptCalls[0].keyID, longFingerprint)
    }

    /// Passphrase stored under long fingerprint is NOT found when the short ID is used for lookup
    func testDecryptWithKeyID_shortIDRecognized_passphraseStoredUnderLongID_missesKeystore() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.privateKeyIDs = [longFingerprint]
        mockPGP.keyIDs = [longFingerprint]
        mockPGP.selectedKeyForPassphrase = shortID

        // Store passphrase under the LONG fingerprint.
        keychain.add(string: "stored-under-long", for: AppKeychain.getPGPKeyPassphraseKey(keyID: longFingerprint))

        _ = try agent.decrypt(encryptedData: testEncryptedData, keyID: shortID, requestPGPKeyPassphrase: passphraseCallback("from-request"))

        // Backend requests passphrase with short ID — keystore lookup misses, falls through to request.
        XCTAssertEqual(mockPGP.resolvedPassphrases, ["from-request"])
        XCTAssertEqual(passphraseRequests, [shortID])
    }

    /// Encrypt with short ID: when containsPublicKey matches (via suffix), the short ID is passed to interface.
    func testEncryptWithKeyID_shortIDRecognized_shortIDFlowsThrough() throws {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.publicKeyIDs = [longFingerprint]
        mockPGP.keyIDs = [longFingerprint]

        _ = try agent.encrypt(plainData: testDecryptedData, keyID: shortID)

        XCTAssertEqual(mockPGP.containsPublicKeyCalls, [shortID])
        XCTAssertEqual(mockPGP.encryptCalls[0].keyID, shortID)
    }

    // MARK: - Encrypt passthrough tests (for completeness of mock interaction)

    /// encrypt(plainData:keyID:) calls containsPublicKey and passes data through.
    func testEncryptWithKeyID_keyFound_callsInterface() throws {
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.publicKeyIDs = [longFingerprint]

        let result = try agent.encrypt(plainData: testDecryptedData, keyID: longFingerprint)

        XCTAssertEqual(result, mockPGP.encryptResult)
        XCTAssertEqual(mockPGP.encryptCalls.count, 1)
        XCTAssertEqual(mockPGP.encryptCalls[0].keyID, longFingerprint)
        XCTAssertEqual(mockPGP.encryptCalls[0].plainData, testDecryptedData)
    }

    /// encrypt with unknown key and single available key falls back.
    func testEncryptWithKeyID_keyNotFound_singleKey_fallsBack() throws {
        let shortID = "e9444483"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.publicKeyIDs = []
        mockPGP.keyIDs = [longFingerprint]

        let result = try agent.encrypt(plainData: testDecryptedData, keyID: shortID)

        XCTAssertEqual(result, mockPGP.encryptResult)
        XCTAssertEqual(mockPGP.containsPublicKeyCalls, [shortID])
        XCTAssertEqual(mockPGP.encryptCalls[0].keyID, longFingerprint)
    }

    /// encrypt with unknown key and multiple keys throws.
    func testEncryptWithKeyID_keyNotFound_multipleKeys_throws() {
        mockPGP.publicKeyIDs = []
        mockPGP.keyIDs = ["4712286271220db299883ea7062e678da1024dae", "787eae1a5fa3e749aa34cc6aa0645ebed862027e"]

        XCTAssertThrowsError(try agent.encrypt(plainData: testDecryptedData, keyID: "a1024dae")) { error in
            XCTAssertEqual(error as? AppError, AppError.pgpPublicKeyNotFound(keyID: "a1024dae"))
        }
        XCTAssertEqual(mockPGP.encryptCalls.count, 0)
    }

    /// encrypt(plainData:) without keyID passes nil to interface.
    func testEncryptNoKeyID_passesNilToInterface() throws {
        let result = try agent.encrypt(plainData: testDecryptedData)

        XCTAssertEqual(result, mockPGP.encryptResult)
        XCTAssertEqual(mockPGP.encryptCalls.count, 1)
        XCTAssertNil(mockPGP.encryptCalls[0].keyID)
    }

    /// encrypt propagates errors from interface.
    func testEncrypt_interfaceThrows_propagatesError() {
        let shortID = "a1024dae"
        let longFingerprint = "4712286271220db299883ea7062e678da1024dae"
        mockPGP.publicKeyIDs = [longFingerprint]
        mockPGP.encryptError = AppError.encryption

        XCTAssertThrowsError(try agent.encrypt(plainData: testDecryptedData, keyID: shortID)) { error in
            XCTAssertEqual(error as? AppError, AppError.encryption)
        }
    }
}