|
1 | 1 | import { expect } from 'chai'; |
2 | 2 | import sinon from 'sinon'; |
3 | | -import FederationProvider from '../../../../../lib/connection/auth/tokenProvider/FederationProvider'; |
| 3 | +import type nodeFetch from 'node-fetch'; |
| 4 | +import FederationProvider, { |
| 5 | + setFederationFetchForTest, |
| 6 | +} from '../../../../../lib/connection/auth/tokenProvider/FederationProvider'; |
4 | 7 | import ITokenProvider from '../../../../../lib/connection/auth/tokenProvider/ITokenProvider'; |
5 | 8 | import Token from '../../../../../lib/connection/auth/tokenProvider/Token'; |
6 | 9 |
|
@@ -68,6 +71,117 @@ describe('FederationProvider', () => { |
68 | 71 | }); |
69 | 72 | }); |
70 | 73 |
|
| 74 | + describe('exchange path', () => { |
| 75 | + // These tests exercise the federation HTTP exchange — the branch |
| 76 | + // taken when the source JWT's issuer doesn't match the Databricks |
| 77 | + // host. The branch contains the AbortController + node-fetch shim |
| 78 | + // typing fix; without coverage here a regression in those mechanics |
| 79 | + // would only surface in production. |
| 80 | + |
| 81 | + afterEach(() => { |
| 82 | + setFederationFetchForTest(); // restore real node-fetch |
| 83 | + }); |
| 84 | + |
| 85 | + // Helper: build a fake node-fetch Response. |
| 86 | + function buildFakeResponse(opts: { |
| 87 | + ok: boolean; |
| 88 | + status?: number; |
| 89 | + statusText?: string; |
| 90 | + body?: unknown; |
| 91 | + text?: string; |
| 92 | + }): nodeFetch.Response { |
| 93 | + return { |
| 94 | + ok: opts.ok, |
| 95 | + status: opts.status ?? (opts.ok ? 200 : 500), |
| 96 | + statusText: opts.statusText ?? '', |
| 97 | + json: async () => opts.body, |
| 98 | + text: async () => opts.text ?? '', |
| 99 | + } as unknown as nodeFetch.Response; |
| 100 | + } |
| 101 | + |
| 102 | + it('should exchange foreign-issued JWT for a Databricks token', async () => { |
| 103 | + const foreignJwt = createJWT({ iss: 'https://idp.example.com' }); |
| 104 | + const baseProvider = new MockTokenProvider(foreignJwt); |
| 105 | + const federationProvider = new FederationProvider(baseProvider, 'my-workspace.cloud.databricks.com'); |
| 106 | + |
| 107 | + const fetchStub = sinon.stub<Parameters<typeof nodeFetch>, ReturnType<typeof nodeFetch>>().resolves( |
| 108 | + buildFakeResponse({ |
| 109 | + ok: true, |
| 110 | + body: { access_token: 'exchanged-databricks-token', token_type: 'Bearer', expires_in: 3600 }, |
| 111 | + }), |
| 112 | + ); |
| 113 | + setFederationFetchForTest(fetchStub as unknown as typeof nodeFetch); |
| 114 | + |
| 115 | + const token = await federationProvider.getToken(); |
| 116 | + |
| 117 | + expect(token.accessToken).to.equal('exchanged-databricks-token'); |
| 118 | + expect(fetchStub.calledOnce).to.be.true; |
| 119 | + |
| 120 | + // The exchange must POST to the Databricks /oidc/v1/token endpoint. |
| 121 | + const [url, init] = fetchStub.firstCall.args; |
| 122 | + expect(String(url)).to.include('my-workspace.cloud.databricks.com'); |
| 123 | + expect(String(url)).to.include('/oidc/v1/token'); |
| 124 | + expect(init!.method).to.equal('POST'); |
| 125 | + |
| 126 | + // Verify the signal propagates an AbortSignal — this is the cast |
| 127 | + // site that TS 5 type-strictness caught. Runtime-wise it must |
| 128 | + // still be a real AbortSignal-shaped object. |
| 129 | + const passedSignal = init!.signal as unknown as AbortSignal; |
| 130 | + expect(passedSignal, 'fetch init.signal must be set').to.exist; |
| 131 | + expect(typeof passedSignal.aborted, 'signal.aborted must be a boolean').to.equal('boolean'); |
| 132 | + expect(passedSignal.aborted).to.be.false; |
| 133 | + }); |
| 134 | + |
| 135 | + it('should propagate abort from the controller to the signal observed by fetch', async () => { |
| 136 | + const foreignJwt = createJWT({ iss: 'https://idp.example.com' }); |
| 137 | + const baseProvider = new MockTokenProvider(foreignJwt); |
| 138 | + const federationProvider = new FederationProvider(baseProvider, 'my-workspace.cloud.databricks.com', { |
| 139 | + returnOriginalTokenOnFailure: false, |
| 140 | + }); |
| 141 | + |
| 142 | + // Capture the signal so we can assert it implements the standard |
| 143 | + // AbortSignal contract. Resolve immediately with success to avoid |
| 144 | + // the 30s real-timeout path; the point is that the signal is wired |
| 145 | + // up, not to exercise the abort end-to-end. |
| 146 | + let capturedSignal: AbortSignal | undefined; |
| 147 | + const fetchStub = sinon |
| 148 | + .stub<Parameters<typeof nodeFetch>, ReturnType<typeof nodeFetch>>() |
| 149 | + .callsFake(async (_url, init) => { |
| 150 | + capturedSignal = init!.signal as unknown as AbortSignal; |
| 151 | + return buildFakeResponse({ |
| 152 | + ok: true, |
| 153 | + body: { access_token: 'tok', token_type: 'Bearer', expires_in: 3600 }, |
| 154 | + }); |
| 155 | + }); |
| 156 | + setFederationFetchForTest(fetchStub as unknown as typeof nodeFetch); |
| 157 | + |
| 158 | + await federationProvider.getToken(); |
| 159 | + |
| 160 | + expect(capturedSignal, 'signal must reach fetch').to.exist; |
| 161 | + // The signal must implement the standard AbortSignal contract. |
| 162 | + expect(typeof capturedSignal!.aborted).to.equal('boolean'); |
| 163 | + expect(typeof capturedSignal!.addEventListener).to.equal('function'); |
| 164 | + }); |
| 165 | + |
| 166 | + it('should fall back to original token when exchange fails (returnOriginalTokenOnFailure default)', async () => { |
| 167 | + const foreignJwt = createJWT({ iss: 'https://idp.example.com' }); |
| 168 | + const baseProvider = new MockTokenProvider(foreignJwt); |
| 169 | + const federationProvider = new FederationProvider(baseProvider, 'my-workspace.cloud.databricks.com'); |
| 170 | + |
| 171 | + const fetchStub = sinon |
| 172 | + .stub<Parameters<typeof nodeFetch>, ReturnType<typeof nodeFetch>>() |
| 173 | + .resolves(buildFakeResponse({ ok: false, status: 400, statusText: 'Bad Request', text: 'invalid_grant' })); |
| 174 | + setFederationFetchForTest(fetchStub as unknown as typeof nodeFetch); |
| 175 | + |
| 176 | + const token = await federationProvider.getToken(); |
| 177 | + |
| 178 | + // Default behavior is to fall back to the original token on failure. |
| 179 | + // Retries kick in for 5xx; 400 is non-retryable so this should fail |
| 180 | + // fast on the first attempt. |
| 181 | + expect(token.accessToken).to.equal(foreignJwt); |
| 182 | + }); |
| 183 | + }); |
| 184 | + |
71 | 185 | describe('getName', () => { |
72 | 186 | it('should return wrapped name', () => { |
73 | 187 | const baseProvider = new MockTokenProvider('token'); |
|
0 commit comments