Peter Steinbergertest: stabilize release precheck041bf4a5 months agoHistory
803 lines · 31.3 KB · Swift
RawDownload
import Foundationimport Testing@testable import CodexBarCore @Suite(.serialized)struct FactoryStatusProbeFetchTests {    @Test    func `keeps stored Factory cookies available when cached header is not logged in`() async throws {        KeychainCacheStore.setTestStoreForTesting(true)        defer { KeychainCacheStore.setTestStoreForTesting(false) }         let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            if url.host == "app.factory.ai",               url.path == "/api/app/auth/me",               request.value(forHTTPHeaderField: "Cookie")?.contains("stale-cache") == true            {                return Self.makeResponse(url: url, body: "{}", statusCode: 401)            }            if url.host == "api.factory.ai", url.path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/organization/subscription/usage" {                let body = """                {                  "usage": {                    "standard": {                      "userTokens": 100,                      "totalAllowance": 1000                    }                  },                  "userId": "user-1"                }                """                return Self.makeResponse(url: url, body: body)            }            return Self.makeResponse(url: url, body: "{}", statusCode: 404)        }         let cookie = try #require(HTTPCookie(properties: [            .domain: "app.factory.ai",            .path: "/",            .name: "session",            .value: "valid-session",        ]))         let sessionFile = try await Self.isolateFactorySessionStore()        defer {            try? FileManager.default.removeItem(at: sessionFile)        }         await FactorySessionStore.shared.clearSession()        CookieHeaderCache.store(provider: .factory, cookieHeader: "session=stale-cache", sourceLabel: "Chrome")        await FactorySessionStore.shared.setCookies([cookie])        await FactorySessionStore.shared.resetInMemoryForTesting()        defer {            CookieHeaderCache.clear(provider: .factory)        }         let probe = FactoryStatusProbe(            timeout: 1.0,            browserDetection: BrowserDetection(                homeDirectory: "/tmp/codexbar-empty-browser-home",                cacheTTL: 0,                fileExists: { _ in false },                directoryContents: { _ in nil }),            transport: FactoryStubTransport())         let snapshot = try await probe.fetch()         #expect(CookieHeaderCache.load(provider: .factory) == nil)        #expect(snapshot.userId == "user-1")        #expect(await FactorySessionStore.shared.getCookies().map(\.value) == ["valid-session"])        #expect(Self.requestTrace() == [            "GET app.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/billing/limits",            "GET api.factory.ai/api/organization/subscription/usage?useCache=true",        ])        await FactorySessionStore.shared.clearSession()    }     @Test    func `preserves stored Factory refresh token when stored cookies are not logged in`() async throws {        KeychainCacheStore.setTestStoreForTesting(true)        defer { KeychainCacheStore.setTestStoreForTesting(false) }         let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            if url.host == "app.factory.ai", url.path == "/api/app/auth/me" {                return Self.makeResponse(url: url, body: "{}", statusCode: 401)            }            if url.host == "api.factory.ai",               url.path == "/api/app/auth/me",               request.value(forHTTPHeaderField: "Cookie")?.contains("stale-session") == true            {                return Self.makeResponse(url: url, body: "{}", statusCode: 401)            }            if url.host == "api.workos.com", url.path == "/user_management/authenticate" {                let requestBody = try Self.requestJSONBody(from: request)                guard requestBody["refresh_token"] as? String == "stale-refresh" else {                    throw URLError(.userAuthenticationRequired)                }                let body = """                {                  "access_token": "fresh-access",                  "refresh_token": "fresh-refresh"                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/organization/subscription/usage" {                let body = """                {                  "usage": {                    "standard": {                      "userTokens": 100,                      "totalAllowance": 1000                    }                  },                  "userId": "user-1"                }                """                return Self.makeResponse(url: url, body: body)            }            return Self.makeResponse(url: url, body: "{}", statusCode: 404)        }         let cookie = try #require(HTTPCookie(properties: [            .domain: "app.factory.ai",            .path: "/",            .name: "session",            .value: "stale-session",        ]))         let sessionFile = try await Self.isolateFactorySessionStore()        defer {            try? FileManager.default.removeItem(at: sessionFile)        }         await FactorySessionStore.shared.clearSession()        CookieHeaderCache.store(provider: .factory, cookieHeader: "session=stale-cache", sourceLabel: "Chrome")        await FactorySessionStore.shared.setCookies([cookie])        await FactorySessionStore.shared.setRefreshToken("stale-refresh")        await FactorySessionStore.shared.resetInMemoryForTesting()        defer {            CookieHeaderCache.clear(provider: .factory)        }         let probe = FactoryStatusProbe(            timeout: 1.0,            browserDetection: BrowserDetection(                homeDirectory: "/tmp/codexbar-empty-browser-home",                cacheTTL: 0,                fileExists: { _ in false },                directoryContents: { _ in nil }),            transport: FactoryStubTransport())         let snapshot = try await probe.fetch()         #expect(CookieHeaderCache.load(provider: .factory) == nil)        #expect(snapshot.userId == "user-1")        #expect(await FactorySessionStore.shared.getCookies().isEmpty)        #expect(await FactorySessionStore.shared.getBearerToken() == "fresh-access")        #expect(await FactorySessionStore.shared.getRefreshToken() == "fresh-refresh")        #expect(Self.requestTrace() == [            "GET app.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/app/auth/me",            "GET app.factory.ai/api/app/auth/me",            "POST api.workos.com/user_management/authenticate",            "GET api.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/billing/limits",            "GET api.factory.ai/api/organization/subscription/usage?useCache=true",        ])        await FactorySessionStore.shared.clearSession()    }     @Test    func `fetches snapshot using cookie header override`() async throws {        let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            let path = url.path            if path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if path == "/api/organization/subscription/usage" {                let body = """                {                  "usage": {                    "startDate": 1700000000000,                    "endDate": 1700003600000,                    "standard": {                      "userTokens": 100,                      "orgTotalTokensUsed": 250,                      "totalAllowance": 1000,                      "usedRatio": 0.10                    },                    "premium": {                      "userTokens": 10,                      "orgTotalTokensUsed": 20,                      "totalAllowance": 100,                      "usedRatio": 0.10                    }                  },                  "userId": "user-1"                }                """                return Self.makeResponse(url: url, body: body)            }            return Self.makeResponse(url: url, body: "{}", statusCode: 404)        }         let probe = FactoryStatusProbe(browserDetection: BrowserDetection(cacheTTL: 0))        let snapshot = try await probe.fetch(cookieHeaderOverride: "access-token=test.jwt.token; session=abc")         #expect(snapshot.standardUserTokens == 100)        #expect(snapshot.standardAllowance == 1000)        #expect(snapshot.standardUsedRatio == 0.10)        #expect(snapshot.premiumUserTokens == 10)        #expect(snapshot.premiumUsedRatio == 0.10)        #expect(snapshot.userId == "user-1")        #expect(snapshot.planName == "Team")         let usage = snapshot.toUsageSnapshot()        #expect(usage.primary?.usedPercent == 10)        #expect(usage.secondary?.usedPercent == 10)    }     @Test    func `uses bearer subject when auth profile omits user id`() async throws {        let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            if url.path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  },                  "userProfile": {                    "email": "user@example.com",                    "role": "member"                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/billing/limits" {                return Self.makeResponse(url: url, body: "{}", statusCode: 404)            }            if url.path == "/api/organization/subscription/usage" {                guard URLComponents(url: url, resolvingAgainstBaseURL: false)?                    .queryItems?                    .contains(where: { $0.name == "userId" && $0.value == "user_jwt" }) == true                else {                    return Self.makeResponse(                        url: url,                        body: #"{"detail":"Must be manager to get usage for other users"}"#,                        statusCode: 403)                }                let body = """                {                  "usage": {                    "standard": {                      "userTokens": 100,                      "totalAllowance": 1000,                      "usedRatio": 0.10                    }                  },                  "userId": "user_jwt"                }                """                return Self.makeResponse(url: url, body: body)            }            return Self.makeResponse(url: url, body: "{}", statusCode: 404)        }         let token = Self.makeJWT(payload: ["sub": "user_jwt"])        let probe = FactoryStatusProbe(browserDetection: BrowserDetection(cacheTTL: 0))        let snapshot = try await probe.fetch(cookieHeaderOverride: "access-token=\(token); session=abc")         #expect(snapshot.userId == "user_jwt")        #expect(snapshot.standardUserTokens == 100)        #expect(Self.requestTrace() == [            "GET app.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/billing/limits",            "GET app.factory.ai/api/organization/subscription/usage?useCache=true&userId=user_jwt",        ])    }     @Test    func `falls back to legacy usage when billing limits request fails`() async throws {        let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            if url.path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/billing/limits" {                throw URLError(.timedOut)            }            if url.path == "/api/organization/subscription/usage" {                let body = """                {                  "usage": {                    "standard": {                      "userTokens": 100,                      "totalAllowance": 1000,                      "usedRatio": 0.10                    },                    "premium": {                      "userTokens": 20,                      "totalAllowance": 100,                      "usedRatio": 0.20                    }                  },                  "userId": "user-1"                }                """                return Self.makeResponse(url: url, body: body)            }            return Self.makeResponse(url: url, body: "{}", statusCode: 404)        }         let probe = FactoryStatusProbe(browserDetection: BrowserDetection(cacheTTL: 0))        let snapshot = try await probe.fetch(cookieHeaderOverride: "access-token=test.jwt.token; session=abc")        let usage = snapshot.toUsageSnapshot()         #expect(snapshot.tokenRateLimits == nil)        #expect(usage.primary?.usedPercent == 10)        #expect(usage.secondary?.usedPercent == 20)        #expect(Self.requestTrace() == [            "GET app.factory.ai/api/app/auth/me",            "GET api.factory.ai/api/billing/limits",            "GET app.factory.ai/api/organization/subscription/usage?useCache=true",        ])    }     @Test    func `falls back to legacy usage when billing limits rejects auth`() async throws {        let registered = URLProtocol.registerClass(FactoryStubURLProtocol.self)        defer {            if registered {                URLProtocol.unregisterClass(FactoryStubURLProtocol.self)            }            FactoryStubURLProtocol.handler = nil            FactoryStubURLProtocol.requests = []        }        FactoryStubURLProtocol.requests = []         FactoryStubURLProtocol.handler = { request in            guard let url = request.url else { throw URLError(.badURL) }            if url.path == "/api/app/auth/me" {                let body = """                {                  "organization": {                    "id": "org_1",                    "name": "Acme",                    "subscription": {                      "factoryTier": "team",                      "orbSubscription": {                        "plan": { "name": "Team", "id": "plan_1" },                        "status": "active"                      }                    }                  }                }                """                return Self.makeResponse(url: url, body: body)            }            if url.host == "api.factory.ai", url.path == "/api/billing/limits" {                return Self.makeResponse(url: url, body: "{}", statusCode: 403)            }            if url.path == "/api/organization/subscription/usage" {                let body = """                {                  "usage": {                    "standard": {                      "userTokens": 100,                      "totalAllowance": 1000,                      "usedRatio": 0.10

Lines 1–500 of 803.