Factor out keychain storage dependency

Sources/UID2/Internal/Storage.swift

@@ -0,0 +1,16 @@
+//
+//  Storage.swift
+//
+
+import Foundation
+
+struct Storage: Sendable {
+    // Load an identity
+    var loadIdentity: @Sendable () async -> (IdentityPackage?)
+
+    // Store an identity
+    var saveIdentity: @Sendable (_ identityPackage: IdentityPackage) async -> Void
+
+    // Clear stored identity
+    var clearIdentity: @Sendable () async -> Void
+}

Sources/UID2/KeychainManager.swift

@@ -5,14 +5,25 @@
 import Foundation
 import Security
 
+extension Storage {
+    static func keychainStorage() -> Storage {
+        let storage = KeychainManager()
+        return .init(
+            loadIdentity: { await storage.loadIdentity() },
+            saveIdentity: { await storage.saveIdentity($0) },
+            clearIdentity: { await storage.clearIdentity() }
+        )
+    }
+}
+
 /// Securely manages data in the Keychain
 actor KeychainManager {
 
     private let attrAccount = "uid2"
 
     private static let attrService = "auth-state"
 
-    func getIdentityFromKeychain() -> IdentityPackage? {
+    func loadIdentity() -> IdentityPackage? {
         let query = query(with: [
             String(kSecReturnData): true
         ])
@@ -28,13 +39,13 @@ actor KeychainManager {
     }
 
     @discardableResult
-    func saveIdentityToKeychain(_ identityPackage: IdentityPackage) -> Bool {
+    func saveIdentity(_ identityPackage: IdentityPackage) -> Bool {
 
         guard let data = try? identityPackage.toData() else {
             return false
         }
 
-        if let _ = getIdentityFromKeychain() {
+        if let _ = loadIdentity() {
             let query = query()
 
             let attributesToUpdate = [String(kSecValueData): data] as CFDictionary
@@ -53,7 +64,7 @@ actor KeychainManager {
     }
 
     @discardableResult
-    func deleteIdentityFromKeychain() -> Bool {
+    func clearIdentity() -> Bool {
         let status: OSStatus = SecItemDelete(query())
 
         return status == errSecSuccess

Sources/UID2/UID2Manager.swift

@@ -85,7 +85,7 @@ public final actor UID2Manager {
     /// UID2Client for Network API  requests
     private let uid2Client: UID2Client
 
-    private let keychainManager = KeychainManager()
+    private let storage: Storage
 
     /// Background Task for Refreshing UID2 Identity
     private var refreshJob: Task<(), Error>?
@@ -128,18 +128,21 @@ public final actor UID2Manager {
                 isLoggingEnabled: isLoggingEnabled,
                 environment: environment
             ),
+            storage: .keychainStorage(),
             sdkVersion: sdkVersion,
             log: log
         )
     }
 
     internal init(
         uid2Client: UID2Client,
+        storage: Storage,
         sdkVersion: (major: Int, minor: Int, patch: Int),
         log: OSLog,
         dateGenerator: DateGenerator = .init { Date() }
     ) {
         self.uid2Client = uid2Client
+        self.storage = storage
         self.sdkVersion = sdkVersion
         self.log = log
         self.dateGenerator = dateGenerator
@@ -171,7 +174,7 @@ public final actor UID2Manager {
     public func resetIdentity() async {
         os_log("Resetting identity", log: log, type: .debug)
         self.state = nil
-        await keychainManager.deleteIdentityFromKeychain()
+        await storage.clearIdentity()
         await checkIdentityExpiration()
         await checkIdentityRefresh()
     }
@@ -242,7 +245,7 @@ public final actor UID2Manager {
     // MARK: - Internal Identity Lifecycle
 
     private func loadStateFromDisk() async {
-        guard let identity = await keychainManager.getIdentityFromKeychain() else {
+        guard let identity = await storage.loadIdentity() else {
             return
         }
         os_log("Restoring previously persisted identity", log: log, type: .debug)
@@ -308,15 +311,15 @@ public final actor UID2Manager {
             os_log("User opt-out detected", log: log, type: .debug)
             self.state = .optout
             let identityPackageOptOut = IdentityPackage(valid: false, errorMessage: "User Opted Out", identity: nil, status: .optOut)
-            await keychainManager.deleteIdentityFromKeychain()
-            await keychainManager.saveIdentityToKeychain(identityPackageOptOut)
+            await storage.clearIdentity()
+            await storage.saveIdentity(identityPackageOptOut)
             return nil
         } else if let identity, status == .established {
             self.state = .established(identity)
             // Not needed for loadFromDisk, but is needed for initial setting of Identity
             let identityPackage = IdentityPackage(valid: true, errorMessage: statusText, identity: identity, status: .established)
             os_log("Updating storage (Status: %@)", log: log, status.debugDescription)
-            await keychainManager.saveIdentityToKeychain(identityPackage)
+            await storage.saveIdentity(identityPackage)
             return identity
         }
 
@@ -331,7 +334,7 @@ public final actor UID2Manager {
         self.state = State(validatedIdentityPackage)
 
         os_log("Updating storage (Status: %@)", log: log, validatedIdentityPackage.status.debugDescription)
-        await keychainManager.saveIdentityToKeychain(validatedIdentityPackage)
+        await storage.saveIdentity(validatedIdentityPackage)
 
         await checkIdentityRefresh()
         await checkIdentityExpiration()

Tests/UID2Tests/UID2ManagerTests.swift

@@ -32,6 +32,7 @@ final class UID2ManagerTests: XCTestCase {
             uid2Client: UID2Client(
                 sdkVersion: "1.0"
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -51,6 +52,7 @@ final class UID2ManagerTests: XCTestCase {
                 sdkVersion: "1.0",
                 cryptoUtil: testCrypto.cryptoUtil
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -87,6 +89,7 @@ final class UID2ManagerTests: XCTestCase {
                 sdkVersion: "1.0",
                 cryptoUtil: testCrypto.cryptoUtil
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -123,6 +126,7 @@ final class UID2ManagerTests: XCTestCase {
                 sdkVersion: "1.0",
                 cryptoUtil: testCrypto.cryptoUtil
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled,
             dateGenerator: .init({ Date(timeIntervalSince1970: 5) })
@@ -164,6 +168,75 @@ final class UID2ManagerTests: XCTestCase {
         XCTAssertEqual(identityStatus, .established)
     }
 
+    // MARK: Identity Restoration
+
+    func testNoneIdentityRestorationFromStorage() async throws {
+        let manager = UID2Manager(
+            uid2Client: UID2Client(
+                sdkVersion: "1.0"
+            ),
+            storage: .null,
+            sdkVersion: (1, 0, 0),
+            log: .disabled
+        )
+        let expectation = XCTestExpectation()
+        await manager.addInitializationListener {
+            let state = await manager.state
+            XCTAssertEqual(state, .none)
+            expectation.fulfill()
+        }
+        await fulfillment(of: [expectation], timeout: 1)
+    }
+
+    func testOptoutIdentityRestorationFromStorage() async throws {
+        let manager = UID2Manager(
+            uid2Client: UID2Client(
+                sdkVersion: "1.0"
+            ),
+            storage: .init(
+                loadIdentity: {
+                    IdentityPackage(valid: true, errorMessage: nil, identity: nil, status: .optOut)
+                },
+                saveIdentity: { _ in },
+                clearIdentity: { }
+            ),
+            sdkVersion: (1, 0, 0),
+            log: .disabled
+        )
+        let expectation = XCTestExpectation()
+        await manager.addInitializationListener {
+            let state = await manager.state
+            XCTAssertEqual(state, .optout)
+            expectation.fulfill()
+        }
+        await fulfillment(of: [expectation], timeout: 1)
+    }
+
+    func testEstablishedIdentityRestorationFromStorage() async throws {
+        let establishedIdentity = UID2Identity.established()
+        let manager = UID2Manager(
+            uid2Client: UID2Client(
+                sdkVersion: "1.0"
+            ),
+            storage: .init(
+                loadIdentity: {
+                    IdentityPackage(valid: true, errorMessage: nil, identity: establishedIdentity, status: .established)
+                },
+                saveIdentity: { _ in },
+                clearIdentity: { }
+            ),
+            sdkVersion: (1, 0, 0),
+            log: .disabled
+        )
+        let expectation = XCTestExpectation()
+        await manager.addInitializationListener {
+            let state = await manager.state
+            XCTAssertEqual(state, .established(establishedIdentity))
+            expectation.fulfill()
+        }
+        await fulfillment(of: [expectation], timeout: 1)
+    }
+
     // MARK: State Observation
 
     @MainActor
@@ -172,6 +245,7 @@ final class UID2ManagerTests: XCTestCase {
             uid2Client: UID2Client(
                 sdkVersion: "1.0"
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -216,6 +290,7 @@ final class UID2ManagerTests: XCTestCase {
             uid2Client: UID2Client(
                 sdkVersion: "1.0"
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -300,3 +375,12 @@ private extension UID2Identity {
         )
     }
 }
+extension Storage {
+    static var null: Self {
+        Storage(
+            loadIdentity: { nil },
+            saveIdentity: { _ in },
+            clearIdentity: {}
+        )
+    }
+}

UID2Prebid/UID2Prebid.xcodeproj/project.pbxproj

@@ -36,6 +36,7 @@
 		BF16EC4C2C5D8D7100B0CA03 /* UID2Manager.State.swift in Sources */ = {isa = PBXBuildFile; fileRef = BF16EC2E2C5D8D7100B0CA03 /* UID2Manager.State.swift */; };
 		BF16EC4D2C5D8D7100B0CA03 /* UID2Manager.swift in Sources */ = {isa = PBXBuildFile; fileRef = BF16EC2F2C5D8D7100B0CA03 /* UID2Manager.swift */; };
 		BF16EC4E2C5D8D7100B0CA03 /* UID2Settings.swift in Sources */ = {isa = PBXBuildFile; fileRef = BF16EC302C5D8D7100B0CA03 /* UID2Settings.swift */; };
+		BF2AA21A2C650C8900634831 /* Storage.swift in Sources */ = {isa = PBXBuildFile; fileRef = BF2AA2192C650C8900634831 /* Storage.swift */; };
 		BFE641292C5CDCB800E241CF /* UID2Prebid.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = BFE641202C5CDCB800E241CF /* UID2Prebid.framework */; };
 		BFE641392C5CDD0600E241CF /* UID2PrebidTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = BFE641382C5CDD0600E241CF /* UID2PrebidTests.swift */; };
 		BFE6413B2C5CDD0B00E241CF /* UID2Prebid.swift in Sources */ = {isa = PBXBuildFile; fileRef = BFE6413A2C5CDD0B00E241CF /* UID2Prebid.swift */; };
@@ -122,6 +123,7 @@
 		BF16EC2E2C5D8D7100B0CA03 /* UID2Manager.State.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = UID2Manager.State.swift; sourceTree = "<group>"; };
 		BF16EC2F2C5D8D7100B0CA03 /* UID2Manager.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = UID2Manager.swift; sourceTree = "<group>"; };
 		BF16EC302C5D8D7100B0CA03 /* UID2Settings.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = UID2Settings.swift; sourceTree = "<group>"; };
+		BF2AA2192C650C8900634831 /* Storage.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Storage.swift; sourceTree = "<group>"; };
 		BFE641202C5CDCB800E241CF /* UID2Prebid.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = UID2Prebid.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		BFE641282C5CDCB800E241CF /* UID2PrebidTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = UID2PrebidTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		BFE641382C5CDD0600E241CF /* UID2PrebidTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = UID2PrebidTests.swift; sourceTree = "<group>"; };
@@ -212,6 +214,7 @@
 			children = (
 				BF16EC172C5D8D7100B0CA03 /* Broadcaster.swift */,
 				BF16EC182C5D8D7100B0CA03 /* Queue.swift */,
+				BF2AA2192C650C8900634831 /* Storage.swift */,
 			);
 			path = Internal;
 			sourceTree = "<group>";
@@ -609,6 +612,7 @@
 				BF16EC4D2C5D8D7100B0CA03 /* UID2Manager.swift in Sources */,
 				BF16EC3A2C5D8D7100B0CA03 /* Queue.swift in Sources */,
 				BF16EC4A2C5D8D7100B0CA03 /* UID2Client.swift in Sources */,
+				BF2AA21A2C650C8900634831 /* Storage.swift in Sources */,
 				BF16EC352C5D8D7100B0CA03 /* UID2Identity.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;

UID2Prebid/UID2PrebidTests/UID2PrebidTests.swift

@@ -31,6 +31,7 @@ final class UID2PrebidTests: XCTestCase {
             uid2Client: UID2Client(
                 sdkVersion: "1.0"
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -74,6 +75,7 @@ final class UID2PrebidTests: XCTestCase {
             uid2Client: UID2Client(
                 sdkVersion: "1.0"
             ),
+            storage: .null,
             sdkVersion: (1, 0, 0),
             log: .disabled
         )
@@ -160,3 +162,13 @@ private extension UID2Identity {
         )
     }
 }
+
+private extension Storage {
+    static var null: Self {
+        Storage(
+            loadIdentity: { nil },
+            saveIdentity: { _ in },
+            clearIdentity: {}
+        )
+    }
+}