Signal-iOS/SignalServiceKit/tests/Storage/PreKeyStoreImplTest.swift
2025-05-05 12:46:46 -05:00

69 lines
2.3 KiB
Swift

//
// Copyright 2023 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
import Foundation
import XCTest
@testable import SignalServiceKit
class PreKeyStoreImplTest: XCTestCase {
private var mockDB: InMemoryDB!
private var aciPreKeyStore: PreKeyStoreImpl!
private var pniPreKeyStore: PreKeyStoreImpl!
override func setUp() {
super.setUp()
mockDB = InMemoryDB()
aciPreKeyStore = PreKeyStoreImpl(for: .aci)
pniPreKeyStore = PreKeyStoreImpl(for: .pni)
}
func testGeneratingAndStoringPreKeys() {
let generatedKeys = mockDB.write { aciPreKeyStore.generatePreKeyRecords(tx: $0) }
XCTAssertEqual(generatedKeys.count, 100)
mockDB.write { tx in
aciPreKeyStore.storePreKeyRecords(generatedKeys, tx: tx)
}
let lastPreKeyRecord = generatedKeys.last!
let firstPreKeyRecord = generatedKeys.first!
mockDB.read { tx in
XCTAssertEqual(
aciPreKeyStore.loadPreKey(lastPreKeyRecord.id, transaction: tx)?.keyPair.publicKey,
lastPreKeyRecord.keyPair.publicKey
)
XCTAssertEqual(
aciPreKeyStore.loadPreKey(firstPreKeyRecord.id, transaction: tx)?.keyPair.publicKey,
firstPreKeyRecord.keyPair.publicKey
)
XCTAssertNil(pniPreKeyStore.loadPreKey(firstPreKeyRecord.id, transaction: tx))
}
}
func testRemovingPreKeys() {
let generatedKeys = mockDB.write { aciPreKeyStore.generatePreKeyRecords(tx: $0) }
XCTAssertEqual(generatedKeys.count, 100)
let lastPreKeyRecord = generatedKeys.last!
let firstPreKeyRecord = generatedKeys.first!
mockDB.write { tx in
aciPreKeyStore.storePreKeyRecords(generatedKeys, tx: tx)
pniPreKeyStore.storePreKeyRecords(generatedKeys, tx: tx)
aciPreKeyStore.removePreKey(lastPreKeyRecord.id, transaction: tx)
}
mockDB.read { tx in
XCTAssertNil(aciPreKeyStore.loadPreKey(lastPreKeyRecord.id, transaction: tx))
XCTAssertNotNil(aciPreKeyStore.loadPreKey(firstPreKeyRecord.id, transaction: tx))
XCTAssertNotNil(pniPreKeyStore.loadPreKey(lastPreKeyRecord.id, transaction: tx))
XCTAssertNotNil(pniPreKeyStore.loadPreKey(firstPreKeyRecord.id, transaction: tx))
}
}
}