Refactor of OAuth support with working tests, still in progress

packages/core/src/callback/index.ts

@@ -16,8 +16,19 @@ export function pTimeout<T>(promise: Promise<T>, timeout: number): Promise<T> {
   })
 }
 
-export function sleep(timeoutInMs: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, timeoutInMs))
+export function sleep(
+  timeoutInMs: number,
+  signal?: AbortSignal
+): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const timeout = setTimeout(resolve, timeoutInMs)
+    if (signal) {
+      signal.addEventListener('abort', () => {
+        clearTimeout(timeout)
+        reject(new DOMException('Aborted', 'AbortError'))
+      })
+    }
+  })
 }
 
 /**

packages/node/src/app/analytics-node.ts

@@ -56,7 +56,7 @@ export class Analytics extends NodeEmitter implements CoreAnalytics {
           typeof settings.httpClient === 'function'
             ? new FetchHTTPClient(settings.httpClient)
             : settings.httpClient ?? new FetchHTTPClient(),
-        oauthSettings: settings.oauthSettings,
+        tokenManagerProps: settings.tokenManagerProps,
       },
       this as NodeEmitter
     )
@@ -73,14 +73,6 @@ export class Analytics extends NodeEmitter implements CoreAnalytics {
     return version
   }
 
-  get oauthSettings() {
-    return this._publisher.oauthSettings
-  }
-
-  set oauthSettings(value) {
-    this._publisher.oauthSettings = value
-  }
-
   /**
    * Call this method to stop collecting new events and flush all existing events.
    * This method also waits for any event method-specific callbacks to be triggered,

packages/node/src/app/settings.ts

@@ -1,6 +1,6 @@
 import { ValidationError } from '@segment/analytics-core'
 import { HTTPClient, HTTPFetchFn } from '../lib/http-client'
-import { OauthSettings } from '../lib/oauth-util'
+import { TokenManagerProps } from '../lib/token-manager'
 
 export interface AnalyticsSettings {
   /**
@@ -44,7 +44,7 @@ export interface AnalyticsSettings {
   /**
    * Set up OAuth2 authentication between the client and Segment's endpoints
    */
-  oauthSettings?: OauthSettings
+  tokenManagerProps?: TokenManagerProps
 }
 
 export const validateSettings = (settings: AnalyticsSettings) => {

packages/node/src/lib/__tests__/oauth-util.test.ts → packages/node/src/lib/__tests__/token-manager.test.ts

@@ -1,7 +1,7 @@
-import { RefreshToken, OauthData, OauthSettings } from '../oauth-util'
+import { sleep } from '@segment/analytics-core'
 import { TestFetchClient } from '../../__tests__/test-helpers/create-test-analytics'
-import { readFileSync } from 'fs'
 import { HTTPResponse } from '../http-client'
+import { TokenManager, TokenManagerProps } from '../token-manager'
 
 const privateKey = Buffer.from(`-----BEGIN PRIVATE KEY-----
 MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDVll7uJaH322IN
@@ -53,68 +53,110 @@ const createOAuthError = (overrides: Partial<HTTPResponse> = {}) => {
   }) as Promise<HTTPResponse>
 }
 
-const getOauthData = () => {
-  const oauthSettings = {
+const getTokenManager = () => {
+  const tokenManagerProps = {
+    httpClient: testClient,
+    maxRetries: 3,
     clientId: 'clientId',
     clientKey: privateKey,
     keyId: 'keyId',
     scope: 'scope',
     authServer: 'http://127.0.0.1:1234',
-  } as OauthSettings
+  } as TokenManagerProps
 
-  const oauthData = {
-    httpClient: testClient,
-    settings: oauthSettings,
-    maxRetries: 3,
-  } as unknown as OauthData
-  return oauthData
+  return new TokenManager(tokenManagerProps)
 }
 
-test('OAuth Success', async () => {
-  fetcher.mockReturnValueOnce(
-    createOAuthSuccess({
-      access_token: 'token',
-      expires_in: '100',
-    })
-  )
-
-  const oauthData = getOauthData()
-
-  RefreshToken(oauthData)
-  await oauthData.refreshPromise
-
-  expect(oauthData.refreshTimer).toBeDefined()
-  expect(oauthData.refreshPromise).toBeUndefined()
-  expect(oauthData.token).toBe('token')
-  expect(fetcher).toHaveBeenCalledTimes(1)
-})
-
-test('OAuth retry failure', async () => {
-  fetcher.mockReturnValue(createOAuthError({ status: 425 }))
-
-  const oauthData = getOauthData()
-
-  RefreshToken(oauthData)
-  await expect(oauthData.refreshPromise).rejects.toThrowError(
-    'Retry limit reached - Foo'
-  )
-
-  expect(oauthData.refreshTimer).toBeUndefined()
-  expect(oauthData.refreshPromise).toBeUndefined()
-  expect(oauthData.token).toBeUndefined()
-  expect(fetcher).toHaveBeenCalledTimes(3)
-})
-
-test('OAuth immediate failure', async () => {
-  fetcher.mockReturnValue(createOAuthError({ status: 400 }))
-
-  const oauthData = getOauthData()
-
-  RefreshToken(oauthData)
-  await expect(oauthData.refreshPromise).rejects.toThrowError('Foo')
-
-  expect(oauthData.refreshTimer).toBeUndefined()
-  expect(oauthData.refreshPromise).toBeUndefined()
-  expect(oauthData.token).toBeUndefined()
-  expect(fetcher).toHaveBeenCalledTimes(1)
-})
+test(
+  'OAuth Success',
+  async () => {
+    fetcher.mockReturnValueOnce(
+      createOAuthSuccess({
+        access_token: 'token',
+        expires_in: 100,
+      })
+    )
+
+    const tokenManager = getTokenManager()
+    const token = await tokenManager.getAccessToken()
+    tokenManager.stopPoller()
+
+    expect(tokenManager.isValidToken(token)).toBeTruthy()
+    expect(token.access_token).toBe('token')
+    expect(token.expires_in).toBe(100)
+    expect(fetcher).toHaveBeenCalledTimes(1)
+  },
+  30 * 1000
+)
+
+test(
+  'OAuth retry failure',
+  async () => {
+    fetcher.mockReturnValue(createOAuthError({ status: 425 }))
+
+    const tokenManager = getTokenManager()
+
+    await expect(tokenManager.getAccessToken()).rejects.toThrowError('Foo')
+    tokenManager.stopPoller()
+
+    expect(fetcher).toHaveBeenCalledTimes(3)
+  },
+  30 * 1000
+)
+
+test(
+  'OAuth immediate failure',
+  async () => {
+    fetcher.mockReturnValue(createOAuthError({ status: 400 }))
+
+    const tokenManager = getTokenManager()
+
+    await expect(tokenManager.getAccessToken()).rejects.toThrowError('Foo')
+    tokenManager.stopPoller()
+
+    expect(fetcher).toHaveBeenCalledTimes(1)
+  },
+  30 * 1000
+)
+
+test(
+  'OAuth rate limit',
+  async () => {
+    fetcher
+      .mockReturnValueOnce(
+        createOAuthError({
+          status: 429,
+          headers: { 'X-RateLimit-Reset': Date.now() + 1000 },
+        })
+      )
+      .mockReturnValueOnce(
+        createOAuthError({
+          status: 429,
+          headers: { 'X-RateLimit-Reset': Date.now() + 1000 },
+        })
+      )
+      .mockReturnValue(
+        createOAuthSuccess({
+          access_token: 'token',
+          expires_in: 100,
+        })
+      )
+
+    const tokenManager = getTokenManager()
+
+    const tokenPromise = tokenManager.getAccessToken()
+    await sleep(250)
+    expect(fetcher).toHaveBeenCalledTimes(1)
+    await sleep(250)
+    expect(fetcher).toHaveBeenCalledTimes(2)
+    await sleep(350)
+    expect(fetcher).toHaveBeenCalledTimes(3)
+
+    const token = await tokenPromise
+    expect(tokenManager.isValidToken(token)).toBeTruthy()
+    expect(token.access_token).toBe('token')
+    expect(token.expires_in).toBe(100)
+    expect(fetcher).toHaveBeenCalledTimes(3)
+  },
+  30 * 1000
+)