mirror of
https://github.com/PostHog/posthog.git
synced 2024-11-24 09:14:46 +01:00
305 lines
11 KiB
TypeScript
305 lines
11 KiB
TypeScript
jest.mock('../../src/utils/now', () => {
|
|
return {
|
|
now: jest.fn(() => Date.now()),
|
|
}
|
|
})
|
|
import { BASE_REDIS_KEY, CELERY_TASK_ID, HogWatcher, HogWatcherState } from '../../src/cdp/hog-watcher'
|
|
import { CdpRedis, createCdpRedisPool } from '../../src/cdp/redis'
|
|
import { HogFunctionInvocationResult } from '../../src/cdp/types'
|
|
import { Hub } from '../../src/types'
|
|
import { closeHub, createHub } from '../../src/utils/db/hub'
|
|
import { delay } from '../../src/utils/utils'
|
|
import { createInvocation } from './fixtures'
|
|
import { deleteKeysWithPrefix } from './helpers/redis'
|
|
|
|
const mockNow: jest.Mock = require('../../src/utils/now').now as any
|
|
|
|
const createResult = (options: {
|
|
id: string
|
|
duration?: number
|
|
finished?: boolean
|
|
error?: string
|
|
}): HogFunctionInvocationResult => {
|
|
return {
|
|
invocation: {
|
|
...createInvocation({ id: options.id }),
|
|
id: 'invocation-id',
|
|
teamId: 2,
|
|
timings: [
|
|
{
|
|
kind: 'async_function',
|
|
duration_ms: options.duration ?? 0,
|
|
},
|
|
],
|
|
},
|
|
finished: options.finished ?? true,
|
|
error: options.error,
|
|
logs: [],
|
|
}
|
|
}
|
|
|
|
describe('HogWatcher', () => {
|
|
describe('integration', () => {
|
|
let now: number
|
|
let hub: Hub
|
|
let watcher: HogWatcher
|
|
let mockCeleryApplyAsync: jest.Mock
|
|
let redis: CdpRedis
|
|
|
|
beforeEach(async () => {
|
|
hub = await createHub()
|
|
hub.celery.applyAsync = mockCeleryApplyAsync = jest.fn()
|
|
|
|
now = 1720000000000
|
|
mockNow.mockReturnValue(now)
|
|
|
|
redis = createCdpRedisPool(hub)
|
|
await deleteKeysWithPrefix(redis, BASE_REDIS_KEY)
|
|
|
|
watcher = new HogWatcher(hub, redis)
|
|
})
|
|
|
|
const advanceTime = (ms: number) => {
|
|
now += ms
|
|
mockNow.mockReturnValue(now)
|
|
}
|
|
|
|
const reallyAdvanceTime = async (ms: number) => {
|
|
advanceTime(ms)
|
|
await delay(ms)
|
|
}
|
|
|
|
afterEach(async () => {
|
|
jest.useRealTimers()
|
|
await closeHub(hub)
|
|
jest.clearAllMocks()
|
|
})
|
|
|
|
it('should retrieve empty state', async () => {
|
|
const res = await watcher.getStates(['id1', 'id2'])
|
|
expect(res).toMatchInlineSnapshot(`
|
|
Object {
|
|
"id1": Object {
|
|
"rating": 1,
|
|
"state": 1,
|
|
"tokens": 10000,
|
|
},
|
|
"id2": Object {
|
|
"rating": 1,
|
|
"state": 1,
|
|
"tokens": 10000,
|
|
},
|
|
}
|
|
`)
|
|
})
|
|
|
|
const cases: [{ cost: number; state: number }, HogFunctionInvocationResult[]][] = [
|
|
[{ cost: 0, state: 1 }, [createResult({ id: 'id1' })]],
|
|
[
|
|
{ cost: 0, state: 1 },
|
|
[createResult({ id: 'id1' }), createResult({ id: 'id1' }), createResult({ id: 'id1' })],
|
|
],
|
|
[
|
|
{ cost: 0, state: 1 },
|
|
[
|
|
createResult({ id: 'id1', duration: 10 }),
|
|
createResult({ id: 'id1', duration: 20 }),
|
|
createResult({ id: 'id1', duration: 100 }),
|
|
],
|
|
],
|
|
[
|
|
{ cost: 12, state: 1 },
|
|
[
|
|
createResult({ id: 'id1', duration: 1000 }),
|
|
createResult({ id: 'id1', duration: 1000 }),
|
|
createResult({ id: 'id1', duration: 1000 }),
|
|
],
|
|
],
|
|
[{ cost: 20, state: 1 }, [createResult({ id: 'id1', duration: 5000 })]],
|
|
[{ cost: 40, state: 1 }, [createResult({ id: 'id1', duration: 10000 })]],
|
|
[
|
|
{ cost: 141, state: 1 },
|
|
[
|
|
createResult({ id: 'id1', duration: 5000 }),
|
|
createResult({ id: 'id1', duration: 10000 }),
|
|
createResult({ id: 'id1', duration: 20000 }),
|
|
],
|
|
],
|
|
|
|
[{ cost: 100, state: 1 }, [createResult({ id: 'id1', error: 'errored!' })]],
|
|
]
|
|
|
|
it.each(cases)('should update tokens based on results %s %s', async (expectedScore, results) => {
|
|
await watcher.observeResults(results)
|
|
const result = await watcher.getState('id1')
|
|
|
|
expect(hub.CDP_WATCHER_BUCKET_SIZE - result.tokens).toEqual(expectedScore.cost)
|
|
expect(result.state).toEqual(expectedScore.state)
|
|
})
|
|
|
|
it('should max out scores', async () => {
|
|
let lotsOfResults = Array(10000).fill(createResult({ id: 'id1', error: 'error!' }))
|
|
|
|
await watcher.observeResults(lotsOfResults)
|
|
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": -0.0001,
|
|
"state": 3,
|
|
"tokens": -1,
|
|
}
|
|
`)
|
|
|
|
lotsOfResults = Array(10000).fill(createResult({ id: 'id2' }))
|
|
|
|
await watcher.observeResults(lotsOfResults)
|
|
|
|
expect(await watcher.getState('id2')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 1,
|
|
"state": 1,
|
|
"tokens": 10000,
|
|
}
|
|
`)
|
|
})
|
|
|
|
it('should refill over time', async () => {
|
|
hub.CDP_WATCHER_REFILL_RATE = 10
|
|
await watcher.observeResults([
|
|
createResult({ id: 'id1', duration: 10000 }),
|
|
createResult({ id: 'id1', duration: 10000 }),
|
|
createResult({ id: 'id1', duration: 10000 }),
|
|
])
|
|
|
|
expect((await watcher.getState('id1')).tokens).toMatchInlineSnapshot(`9880`)
|
|
advanceTime(1000)
|
|
expect((await watcher.getState('id1')).tokens).toMatchInlineSnapshot(`9890`)
|
|
advanceTime(10000)
|
|
expect((await watcher.getState('id1')).tokens).toMatchInlineSnapshot(`9990`)
|
|
})
|
|
|
|
it('should remain disabled for period', async () => {
|
|
const badResults = Array(100).fill(createResult({ id: 'id1', error: 'error!' }))
|
|
|
|
await watcher.observeResults(badResults)
|
|
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledTimes(1)
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledWith(CELERY_TASK_ID, [
|
|
'id1',
|
|
HogWatcherState.disabledForPeriod,
|
|
])
|
|
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 0,
|
|
"state": 3,
|
|
"tokens": 0,
|
|
}
|
|
`)
|
|
|
|
advanceTime(10000)
|
|
|
|
// Should still be disabled even though tokens have been refilled
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 0.01,
|
|
"state": 3,
|
|
"tokens": 100,
|
|
}
|
|
`)
|
|
})
|
|
|
|
describe('forceStateChange', () => {
|
|
it('should force healthy', async () => {
|
|
await watcher.forceStateChange('id1', HogWatcherState.healthy)
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 1,
|
|
"state": 1,
|
|
"tokens": 10000,
|
|
}
|
|
`)
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledWith(CELERY_TASK_ID, ['id1', HogWatcherState.healthy])
|
|
})
|
|
it('should force degraded', async () => {
|
|
await watcher.forceStateChange('id1', HogWatcherState.degraded)
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 0.8,
|
|
"state": 1,
|
|
"tokens": 8000,
|
|
}
|
|
`)
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledWith(CELERY_TASK_ID, ['id1', HogWatcherState.degraded])
|
|
})
|
|
it('should force disabledForPeriod', async () => {
|
|
await watcher.forceStateChange('id1', HogWatcherState.disabledForPeriod)
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 0,
|
|
"state": 3,
|
|
"tokens": 0,
|
|
}
|
|
`)
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledWith(CELERY_TASK_ID, [
|
|
'id1',
|
|
HogWatcherState.disabledForPeriod,
|
|
])
|
|
})
|
|
it('should force disabledIndefinitely', async () => {
|
|
await watcher.forceStateChange('id1', HogWatcherState.disabledIndefinitely)
|
|
expect(await watcher.getState('id1')).toMatchInlineSnapshot(`
|
|
Object {
|
|
"rating": 0,
|
|
"state": 4,
|
|
"tokens": 0,
|
|
}
|
|
`)
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledWith(CELERY_TASK_ID, [
|
|
'id1',
|
|
HogWatcherState.disabledIndefinitely,
|
|
])
|
|
})
|
|
})
|
|
|
|
describe('disable logic', () => {
|
|
beforeEach(() => {
|
|
hub.CDP_WATCHER_BUCKET_SIZE = 100
|
|
hub.CDP_WATCHER_DISABLED_TEMPORARY_TTL = 1 // Shorter ttl to help with testing
|
|
hub.CDP_WATCHER_DISABLED_TEMPORARY_MAX_COUNT = 3
|
|
})
|
|
|
|
it('count the number of times it has been disabled', async () => {
|
|
// Trigger the temporary disabled state 3 times
|
|
for (let i = 0; i < 2; i++) {
|
|
await watcher.observeResults([createResult({ id: 'id1', error: 'error!' })])
|
|
expect((await watcher.getState('id1')).state).toEqual(HogWatcherState.disabledForPeriod)
|
|
await reallyAdvanceTime(1000)
|
|
expect((await watcher.getState('id1')).state).toEqual(HogWatcherState.degraded)
|
|
}
|
|
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledTimes(2)
|
|
expect(mockCeleryApplyAsync.mock.calls[0]).toEqual([
|
|
CELERY_TASK_ID,
|
|
['id1', HogWatcherState.disabledForPeriod],
|
|
])
|
|
expect(mockCeleryApplyAsync.mock.calls[1]).toEqual([
|
|
CELERY_TASK_ID,
|
|
['id1', HogWatcherState.disabledForPeriod],
|
|
])
|
|
|
|
await watcher.observeResults([createResult({ id: 'id1', error: 'error!' })])
|
|
expect((await watcher.getState('id1')).state).toEqual(HogWatcherState.disabledIndefinitely)
|
|
await reallyAdvanceTime(1000)
|
|
expect((await watcher.getState('id1')).state).toEqual(HogWatcherState.disabledIndefinitely)
|
|
|
|
expect(mockCeleryApplyAsync).toHaveBeenCalledTimes(3)
|
|
expect(mockCeleryApplyAsync.mock.calls[2]).toEqual([
|
|
CELERY_TASK_ID,
|
|
['id1', HogWatcherState.disabledIndefinitely],
|
|
])
|
|
})
|
|
})
|
|
})
|
|
})
|