From f6586537b6fc8def98ccc6a25e8d707f4fe33c28 Mon Sep 17 00:00:00 2001
From: Cyrus Ho <hello@cyrusho.me>
Date: Tue, 24 Feb 2026 03:42:32 +0800
Subject: [PATCH] test(sso): add tests for ensureRuntimeDiscovery function

---
 packages/sso/src/oidc/discovery.test.ts | 83 ++++++++++++++++++++++++-
 1 file changed, 82 insertions(+), 1 deletion(-)

diff --git a/packages/sso/src/oidc/discovery.test.ts b/packages/sso/src/oidc/discovery.test.ts
index 2186f900cb..5ed24e67ff 100644
--- a/packages/sso/src/oidc/discovery.test.ts
+++ b/packages/sso/src/oidc/discovery.test.ts
@@ -1,7 +1,8 @@
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import {
 	computeDiscoveryUrl,
 	discoverOIDCConfig,
+	ensureRuntimeDiscovery,
 	fetchDiscoveryDocument,
 	needsRuntimeDiscovery,
 	normalizeDiscoveryUrls,
@@ -1150,3 +1151,83 @@ describe("OIDC Discovery", () => {
 		});
 	});
 });
+
+describe("ensureRuntimeDiscovery", () => {
+	const isTrustedOrigin = () => true;
+	const issuer = "https://idp.example.com";
+
+	const baseConfig = {
+		issuer,
+		clientId: "client-id",
+		clientSecret: "client-secret",
+		pkce: true,
+		discoveryEndpoint: `${issuer}/.well-known/openid-configuration`,
+	};
+
+	const mockDiscoveryDoc = createMockDiscoveryDocument();
+
+	beforeEach(() => {
+		vi.mocked(betterFetch).mockResolvedValue({
+			data: mockDiscoveryDoc,
+			error: null,
+		});
+	});
+
+	it("returns config unchanged when discovery is not needed", async () => {
+		const completeConfig = {
+			...baseConfig,
+			tokenEndpoint: `${issuer}/oauth2/token`,
+			jwksEndpoint: `${issuer}/.well-known/jwks.json`,
+		};
+		const result = await ensureRuntimeDiscovery(
+			completeConfig,
+			issuer,
+			isTrustedOrigin,
+		);
+		expect(result).toBe(completeConfig);
+		expect(betterFetch).not.toHaveBeenCalled();
+	});
+
+	it("hydrates missing endpoints when discovery is needed", async () => {
+		const result = await ensureRuntimeDiscovery(
+			baseConfig,
+			issuer,
+			isTrustedOrigin,
+		);
+		expect(result.authorizationEndpoint).toBe(
+			mockDiscoveryDoc.authorization_endpoint,
+		);
+		expect(result.tokenEndpoint).toBe(mockDiscoveryDoc.token_endpoint);
+		expect(result.jwksEndpoint).toBe(mockDiscoveryDoc.jwks_uri);
+		expect(result.userInfoEndpoint).toBe(mockDiscoveryDoc.userinfo_endpoint);
+	});
+
+	it("preserves existing config fields not returned by discovery", async () => {
+		const result = await ensureRuntimeDiscovery(
+			baseConfig,
+			issuer,
+			isTrustedOrigin,
+		);
+		expect(result.clientId).toBe(baseConfig.clientId);
+		expect(result.clientSecret).toBe(baseConfig.clientSecret);
+		expect(result.pkce).toBe(baseConfig.pkce);
+	});
+
+	it("throws when discovery fails", async () => {
+		vi.mocked(betterFetch).mockResolvedValue({
+			data: null,
+			error: { message: "Network error" },
+		});
+		await expect(
+			ensureRuntimeDiscovery(baseConfig, issuer, isTrustedOrigin),
+		).rejects.toThrow(DiscoveryError);
+	});
+
+	it("throws DiscoveryError for untrusted origin", async () => {
+		await expect(
+			ensureRuntimeDiscovery(baseConfig, issuer, () => false),
+		).rejects.toThrow(
+			expect.objectContaining({ code: "discovery_untrusted_origin" }),
+		);
+	});
+});