From a3942db823eca1c457680a2fe00d130d38e18ebd Mon Sep 17 00:00:00 2001 From: Tiina Turban Date: Wed, 18 May 2022 17:29:01 +0200 Subject: [PATCH 1/5] refactor(plugin-server): yeet redis queue --- .../src/main/ingestion-queues/celery-queue.ts | 226 ------------------ .../src/main/ingestion-queues/queue.ts | 74 +----- plugin-server/tests/queue.test.ts | 50 ++++ 3 files changed, 53 insertions(+), 297 deletions(-) delete mode 100644 plugin-server/src/main/ingestion-queues/celery-queue.ts create mode 100644 plugin-server/tests/queue.test.ts diff --git a/plugin-server/src/main/ingestion-queues/celery-queue.ts b/plugin-server/src/main/ingestion-queues/celery-queue.ts deleted file mode 100644 index dadfa05b8f2a..000000000000 --- a/plugin-server/src/main/ingestion-queues/celery-queue.ts +++ /dev/null @@ -1,226 +0,0 @@ -import { Queue } from '../../types' -import { Base } from '../../utils/celery/base' -import { Message } from '../../utils/celery/message' -import { status } from '../../utils/status' - -type Handler = (...args: any[]) => Promise - -export class CeleryQueue extends Base implements Queue { - handlers: Record = {} - activeTasks: Set> = new Set() - - /** - * register task handler on worker handlers - * @method Worker#register - * @param {String} name the name of task for dispatching. - * @param {Function} handler the function for task handling - * - * @example - * worker.register('tasks.add', (a, b) => a + b); - * worker.start(); - */ - public register(name: string, handler: Handler): void { - if (!handler) { - throw new Error('Undefined handler') - } - if (this.handlers[name]) { - throw new Error('Already handler setted') - } - - this.handlers[name] = function registHandler(...args: any[]): Promise { - try { - return Promise.resolve(handler(...args)) - } catch (err) { - return Promise.reject(err) - } - } - } - - /** - * start celery worker to run - * @method Worker#start - * @example - * worker.register('tasks.add', (a, b) => a + b); - * worker.start(); - */ - public start(): Promise { - status.info('🍆', 'Starting Celery worker...') - return this.run().catch((error) => - status.error('⚠️', 'An error occured while starting Celery worker:\n', error) - ) - } - - /** - * Pause the worker. await the response to be sure all pending `processNextTick` events have finished. - * @method Worker#pause - */ - public pause(): Promise { - return this.broker.pause() - } - - /** - * Resume the worker - * @method Worker#pause - */ - public resume(): void { - this.broker.resume() - } - - /** - * Is the worker paused - * @method Worker#isPaused - * - * @returns {boolean} - */ - public isPaused(): boolean { - return this.broker.isPaused() - } - - /** - * @method Worker#run - * @private - * - * @returns {Promise} - */ - private run(): Promise { - return this.processTasks() - } - - /** - * @method Worker#processTasks - * @private - * - * @returns function results - */ - private processTasks(): Promise { - const consumer = this.getConsumer(this.conf.CELERY_QUEUE) - return consumer() - } - - /** - * @method Worker#getConsumer - * @private - * - * @param {String} queue queue name for task route - */ - private getConsumer(queue: string): () => any { - const onMessage = this.createTaskHandler() - - return () => { - const result = this.broker.subscribe(queue, onMessage) - status.info('✅', `Celery worker subscribed to ${Object.keys(this.handlers).join(', ')}!`) - return result - } - } - - public createTaskHandler(): (message: Message) => any { - const onTaskReceived = (message: Message): any => { - if (!message) { - return Promise.resolve() - } - - let payload: Record | null = null - let taskName = message.headers['task'] - if (!taskName) { - // protocol v1 - payload = message.decode() - taskName = payload['task'] - } - - // strategy - let body - let headers - if (payload == null && !('args' in message.decode())) { - body = message.decode() // message.body; - headers = message.headers - } else if (payload) { - const args = payload['args'] || [] - const kwargs = payload['kwargs'] || {} - const embed = { - callbacks: payload['callbacks'], - errbacks: payload['errbacks'], - chord: payload['chord'], - chain: null, - } - - body = [args, kwargs, embed] - headers = { - lang: payload['lang'], - task: payload['task'], - id: payload['id'], - rootId: payload['root_id'], - parantId: payload['parentId'], - group: payload['group'], - meth: payload['meth'], - shadow: payload['shadow'], - eta: payload['eta'], - expires: payload['expires'], - retries: payload['retries'] || 0, - timelimit: payload['timelimit'] || [null, null], - kwargsrepr: payload['kwargsrepr'], - origin: payload['origin'], - } - } - - // request - const [args, kwargs /*, embed */] = body as [Array, Record, Record] - const taskId = headers ? headers['id'] : null - - const handler = this.handlers[taskName] - if (!handler) { - throw new Error(`Missing process handler for task ${taskName}`) - } - - console.debug( - `celery.node Received task: ${taskName}[${taskId}], args: ${args}, kwargs: ${JSON.stringify(kwargs)}` - ) - - const timeStart = process.hrtime() - const taskPromise = handler(...args, kwargs).then((result) => { - const diff = process.hrtime(timeStart) - console.debug( - `celery.node Task ${taskName}[${taskId}] succeeded in ${diff[0] + diff[1] / 1e9}s: ${result}` - ) - this.activeTasks.delete(taskPromise) - }) - - // record the executing task - this.activeTasks.add(taskPromise) - - return taskPromise - } - - return onTaskReceived - } - - /** - * @method Worker#whenCurrentJobsFinished - * - * @returns Promise that resolves when all jobs are finished - */ - public async whenCurrentJobsFinished(): Promise { - return Promise.all(Array.from(this.activeTasks)) - } - - /** - * @method Worker#stop - * - * @todo implement here - */ - public async stop(): Promise { - const taskCount = this.activeTasks.size - if (taskCount > 0) { - status.info( - '⌛', - `${taskCount} ${taskCount === 1 ? 'task' : 'tasks'} in progress, waiting for ${ - taskCount === 1 ? 'it' : 'them' - } to finish before disconnecting Celery...` - ) - await this.whenCurrentJobsFinished() - } else { - status.info('👍', 'No tasks in progress, disconnecting Celery...') - } - await this.disconnect() - status.info('🛑', 'Celery worker disconnected!') - } -} diff --git a/plugin-server/src/main/ingestion-queues/queue.ts b/plugin-server/src/main/ingestion-queues/queue.ts index bd9691546763..6bfe858fd1a3 100644 --- a/plugin-server/src/main/ingestion-queues/queue.ts +++ b/plugin-server/src/main/ingestion-queues/queue.ts @@ -1,26 +1,13 @@ import Piscina from '@posthog/piscina' import { PluginEvent, ProcessedPluginEvent } from '@posthog/plugin-scaffold' -import * as Sentry from '@sentry/node' -import { - CeleryTriggeredJobOperation, - Hub, - PluginConfig, - PreIngestionEvent, - Queue, - Team, - WorkerMethods, -} from '../../types' +import { Hub, PreIngestionEvent, Queue, WorkerMethods } from '../../types' import { status } from '../../utils/status' -import { sanitizeEvent, UUIDT } from '../../utils/utils' import { Action } from './../../types' -import { CeleryQueue } from './celery-queue' -import { ingestEvent } from './ingest-event' import { KafkaQueue } from './kafka-queue' interface Queues { ingestion: Queue | null - auxiliary: Queue | null } export function pauseQueueIfWorkerFull( @@ -78,13 +65,8 @@ export async function startQueues( } try { - const redisQueue = startQueueRedis(server, piscina, mergedWorkerMethods) - const queues = { - ingestion: redisQueue, - auxiliary: redisQueue, - } - if (server.KAFKA_ENABLED) { - queues.ingestion = await startQueueKafka(server, mergedWorkerMethods) + const queues: Queues = { + ingestion: await startQueueKafka(server, mergedWorkerMethods), } return queues } catch (error) { @@ -93,56 +75,6 @@ export async function startQueues( } } -function startQueueRedis(server: Hub, piscina: Piscina, workerMethods: WorkerMethods): Queue | null { - const celeryQueue = new CeleryQueue(server.db, server.PLUGINS_CELERY_QUEUE) - - let startQueue = false - - // if kafka is enabled, we'll process events from there - if (!server.KAFKA_ENABLED && server.capabilities.ingestion) { - startQueue = true - - celeryQueue.register( - 'posthog.tasks.process_event.process_event_with_plugins', - async ( - distinct_id: string, - ip: string | null, - site_url: string, - data: Record, - team_id: number, - now: string, - sent_at?: string - ) => { - const event = sanitizeEvent({ - distinct_id, - ip, - site_url, - team_id, - now, - sent_at, - uuid: new UUIDT().toString(), - ...data, - } as PluginEvent) - try { - const checkAndPause = () => pauseQueueIfWorkerFull(() => celeryQueue.pause(), server, piscina) - await ingestEvent(server, workerMethods, event, checkAndPause) - } catch (e) { - Sentry.captureException(e) - } - } - ) - } - - if (startQueue) { - // run in the background - void celeryQueue.start() - - return celeryQueue - } else { - return null - } -} - async function startQueueKafka(server: Hub, workerMethods: WorkerMethods): Promise { if (!server.capabilities.ingestion) { return null diff --git a/plugin-server/tests/queue.test.ts b/plugin-server/tests/queue.test.ts new file mode 100644 index 000000000000..64f5f9665fec --- /dev/null +++ b/plugin-server/tests/queue.test.ts @@ -0,0 +1,50 @@ +import Piscina from '@posthog/piscina' + +import { KafkaQueue } from '../src/main/ingestion-queues/kafka-queue' +import { startQueues } from '../src/main/ingestion-queues/queue' +import { Hub, LogLevel } from '../src/types' +import { createHub } from '../src/utils/db/hub' + +jest.setTimeout(60000) // 60 sec timeout +jest.mock('../src/main/ingestion-queues/kafka-queue') +jest.mock('../src/utils/status') +jest.mock('../src/main/ingestion-queues/ingest-event') + +describe('queue', () => { + describe('capabilities', () => { + let hub: Hub + let piscina: Piscina + let closeHub: () => Promise + + beforeEach(async () => { + ;[hub, closeHub] = await createHub({ + LOG_LEVEL: LogLevel.Warn, + KAFKA_ENABLED: true, + KAFKA_HOSTS: process.env.KAFKA_HOSTS || 'kafka:9092', + }) + piscina = { run: jest.fn() } as any + }) + + afterEach(async () => { + await closeHub() + }) + + it('starts ingestion queue by default', async () => { + const queues = await startQueues(hub, piscina) + + expect(queues).toEqual({ + ingestion: expect.any(KafkaQueue), + }) + }) + + it('handles ingestion being turned off', async () => { + hub.capabilities.ingestion = false + + const queues = await startQueues(hub, piscina) + + expect(queues).toEqual({ + ingestion: null, + }) + }) + }) +}) From 63f40ec99fcbc872877109921938a4b1a0606394 Mon Sep 17 00:00:00 2001 From: Tiina Turban Date: Thu, 19 May 2022 18:39:34 +0200 Subject: [PATCH 2/5] fix --- plugin-server/src/main/pluginsServer.ts | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/plugin-server/src/main/pluginsServer.ts b/plugin-server/src/main/pluginsServer.ts index 81f0e0590d16..f80bb61ad475 100644 --- a/plugin-server/src/main/pluginsServer.ts +++ b/plugin-server/src/main/pluginsServer.ts @@ -61,7 +61,6 @@ export async function startPluginsServer( let hub: Hub | undefined let piscina: Piscina | undefined let queue: Queue | undefined | null // ingestion queue - let redisQueueForPluginJobs: Queue | undefined | null let healthCheckConsumer: Consumer | undefined let jobQueueConsumer: JobQueueConsumerControl | undefined let closeHub: () => Promise | undefined @@ -87,7 +86,6 @@ export async function startPluginsServer( lastActivityCheck && clearInterval(lastActivityCheck) cancelAllScheduledJobs() await queue?.stop() - await redisQueueForPluginJobs?.stop() await pubSub?.stop() await jobQueueConsumer?.stop() await pluginScheduleControl?.stopSchedule() @@ -160,16 +158,10 @@ export async function startPluginsServer( const queues = await startQueues(hub, piscina) - // `queue` refers to the ingestion queue. With Celery ingestion, we only - // have one queue for plugin jobs and ingestion. With Kafka ingestion, we - // use Kafka for events but still start Redis for plugin jobs. - // Thus, if Kafka is disabled, we don't need to call anything on - // redisQueueForPluginJobs, as that will also be the ingestion queue. + // `queue` refers to the ingestion queue. queue = queues.ingestion - redisQueueForPluginJobs = config.KAFKA_ENABLED ? queues.auxiliary : null piscina.on('drain', () => { void queue?.resume() - void redisQueueForPluginJobs?.resume() void jobQueueConsumer?.resume() }) From c046ca63d40d8193f4ccb02eb495857ce9366bdb Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Mon, 23 May 2022 15:24:17 +0200 Subject: [PATCH 3/5] Remove unused imports --- plugin-server/tests/clickhouse/postgres-parity.test.ts | 4 ---- plugin-server/tests/shared/vm.test.ts | 1 - 2 files changed, 5 deletions(-) diff --git a/plugin-server/tests/clickhouse/postgres-parity.test.ts b/plugin-server/tests/clickhouse/postgres-parity.test.ts index 67427e02b2d4..8b2fa94a788a 100644 --- a/plugin-server/tests/clickhouse/postgres-parity.test.ts +++ b/plugin-server/tests/clickhouse/postgres-parity.test.ts @@ -12,10 +12,8 @@ import { } from '../../src/types' import { castTimestampOrNow, UUIDT } from '../../src/utils/utils' import { makePiscina } from '../../src/worker/piscina' -import { createPosthog, DummyPostHog } from '../../src/worker/vm/extensions/posthog' import { delayUntilEventIngested, resetTestDatabaseClickhouse } from '../helpers/clickhouse' import { resetKafka } from '../helpers/kafka' -import { pluginConfig39 } from '../helpers/plugins' import { getFirstTeam, resetTestDatabase } from '../helpers/sql' jest.mock('../../src/utils/status') @@ -29,7 +27,6 @@ const extraServerConfig: Partial = { describe('postgres parity', () => { let hub: Hub let stopServer: () => Promise - let posthog: DummyPostHog let team: Team beforeAll(async () => { @@ -48,7 +45,6 @@ describe('postgres parity', () => { const startResponse = await startPluginsServer(extraServerConfig, makePiscina) hub = startResponse.hub stopServer = startResponse.stop - posthog = createPosthog(hub, pluginConfig39) team = await getFirstTeam(hub) }) diff --git a/plugin-server/tests/shared/vm.test.ts b/plugin-server/tests/shared/vm.test.ts index e8c1acea63f0..e5058604f750 100644 --- a/plugin-server/tests/shared/vm.test.ts +++ b/plugin-server/tests/shared/vm.test.ts @@ -6,7 +6,6 @@ import { JobQueueManager } from '../../src/main/job-queues/job-queue-manager' import { Hub } from '../../src/types' import { PluginConfig, PluginConfigVMResponse } from '../../src/types' import { createHub } from '../../src/utils/db/hub' -import { KafkaProducerWrapper } from '../../src/utils/db/kafka-producer-wrapper' import { delay } from '../../src/utils/utils' import { MAXIMUM_RETRIES } from '../../src/worker/vm/upgrades/export-events' import { createPluginConfigVM } from '../../src/worker/vm/vm' From 85b87d5ea99a036d44ab6d2e5bfad9001bcf9c07 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Mon, 23 May 2022 20:21:45 +0200 Subject: [PATCH 4/5] Adjust tests to new reality --- plugin-server/tests/errors.test.ts | 23 +++++++++++----------- plugin-server/tests/postgres/queue.test.ts | 2 -- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/plugin-server/tests/errors.test.ts b/plugin-server/tests/errors.test.ts index 4ccad18e025a..37fc753e9004 100644 --- a/plugin-server/tests/errors.test.ts +++ b/plugin-server/tests/errors.test.ts @@ -2,7 +2,8 @@ import { startPluginsServer } from '../src/main/pluginsServer' import { Hub, LogLevel, PluginsServerConfig } from '../src/types' import { makePiscina } from '../src/worker/piscina' import { createPosthog, DummyPostHog } from '../src/worker/vm/extensions/posthog' -import { delayUntilEventIngested } from './helpers/clickhouse' +import { delayUntilEventIngested, resetTestDatabaseClickhouse } from './helpers/clickhouse' +import { resetKafka } from './helpers/kafka' import { pluginConfig39 } from './helpers/plugins' import { getErrorForPluginConfig, resetTestDatabase } from './helpers/sql' @@ -10,10 +11,7 @@ jest.setTimeout(60000) // 60 sec timeout const extraServerConfig: Partial = { WORKER_CONCURRENCY: 2, - PLUGINS_CELERY_QUEUE: 'test-plugins-celery-queue-errors', - CELERY_DEFAULT_QUEUE: 'test-celery-default-queue-errors', LOG_LEVEL: LogLevel.Log, - KAFKA_ENABLED: false, } describe('error do not take down ingestion', () => { @@ -21,8 +19,13 @@ describe('error do not take down ingestion', () => { let stopServer: () => Promise let posthog: DummyPostHog + beforeAll(async () => { + await resetKafka() + }) + beforeEach(async () => { - await resetTestDatabase(` + await resetTestDatabase( + ` export async function processEvent (event, { jobs }) { if (event.properties.crash === 'throw') { throw new Error('error thrown in plugin') @@ -33,16 +36,14 @@ describe('error do not take down ingestion', () => { } return event } - `) + `, + extraServerConfig + ) + await resetTestDatabaseClickhouse(extraServerConfig) const startResponse = await startPluginsServer(extraServerConfig, makePiscina) hub = startResponse.hub stopServer = startResponse.stop posthog = createPosthog(hub, pluginConfig39) - - const redis = await hub.redisPool.acquire() - await redis.del(hub.PLUGINS_CELERY_QUEUE) - await redis.del(hub.CELERY_DEFAULT_QUEUE) - await hub.redisPool.release(redis) }) afterEach(async () => { diff --git a/plugin-server/tests/postgres/queue.test.ts b/plugin-server/tests/postgres/queue.test.ts index a9ce8a3fc00c..763924fcd3e0 100644 --- a/plugin-server/tests/postgres/queue.test.ts +++ b/plugin-server/tests/postgres/queue.test.ts @@ -108,7 +108,6 @@ describe('queue', () => { expect(queues).toEqual({ ingestion: expect.any(KafkaQueue), - auxiliary: null, }) }) @@ -119,7 +118,6 @@ describe('queue', () => { expect(queues).toEqual({ ingestion: null, - auxiliary: null, }) }) }) From 4623f0b67357955a2f988a2bc2dbf91918708d9d Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Mon, 23 May 2022 20:40:32 +0200 Subject: [PATCH 5/5] Refactor `queue.test.ts` away --- .../tests/clickhouse/kafka-queue.test.ts | 4 +- plugin-server/tests/postgres/queue.test.ts | 124 ------------------ .../tests/shared/capabilities.test.ts | 45 +++++++ 3 files changed, 46 insertions(+), 127 deletions(-) delete mode 100644 plugin-server/tests/postgres/queue.test.ts create mode 100644 plugin-server/tests/shared/capabilities.test.ts diff --git a/plugin-server/tests/clickhouse/kafka-queue.test.ts b/plugin-server/tests/clickhouse/kafka-queue.test.ts index 77a790346f8c..760020c88e07 100644 --- a/plugin-server/tests/clickhouse/kafka-queue.test.ts +++ b/plugin-server/tests/clickhouse/kafka-queue.test.ts @@ -1,12 +1,10 @@ import Piscina from '@posthog/piscina' -import { ONE_HOUR } from '../../src/config/constants' import { KAFKA_EVENTS_PLUGIN_INGESTION } from '../../src/config/kafka-topics' import { ServerInstance, startPluginsServer } from '../../src/main/pluginsServer' import { LogLevel, PluginsServerConfig } from '../../src/types' import { Hub } from '../../src/types' -import { Client } from '../../src/utils/celery/client' -import { delay, UUIDT } from '../../src/utils/utils' +import { UUIDT } from '../../src/utils/utils' import { makePiscina } from '../../src/worker/piscina' import { createPosthog, DummyPostHog } from '../../src/worker/vm/extensions/posthog' import { writeToFile } from '../../src/worker/vm/extensions/test-utils' diff --git a/plugin-server/tests/postgres/queue.test.ts b/plugin-server/tests/postgres/queue.test.ts deleted file mode 100644 index 763924fcd3e0..000000000000 --- a/plugin-server/tests/postgres/queue.test.ts +++ /dev/null @@ -1,124 +0,0 @@ -import Piscina from '@posthog/piscina' - -import { CeleryQueue } from '../../src/main/ingestion-queues/celery-queue' -import { ingestEvent } from '../../src/main/ingestion-queues/ingest-event' -import { KafkaQueue } from '../../src/main/ingestion-queues/kafka-queue' -import { startQueues } from '../../src/main/ingestion-queues/queue' -import { Hub, LogLevel } from '../../src/types' -import { Client } from '../../src/utils/celery/client' -import { createHub } from '../../src/utils/db/hub' -import { delay } from '../../src/utils/utils' - -jest.setTimeout(60000) // 60 sec timeout -jest.mock('../../src/main/ingestion-queues/kafka-queue') -jest.mock('../../src/utils/status') -jest.mock('../../src/main/ingestion-queues/ingest-event') - -function advanceOneTick() { - return new Promise((resolve) => process.nextTick(resolve)) -} - -async function createTestHub(): Promise<[Hub, () => Promise]> { - const [hub, closeHub] = await createHub({ - REDIS_POOL_MIN_SIZE: 3, - REDIS_POOL_MAX_SIZE: 3, - PLUGINS_CELERY_QUEUE: 'ttt-test-plugins-celery-queue', - CELERY_DEFAULT_QUEUE: 'ttt-test-celery-default-queue', - LOG_LEVEL: LogLevel.Log, - KAFKA_ENABLED: false, - }) - - const redis = await hub.redisPool.acquire() - await redis.del(hub.PLUGINS_CELERY_QUEUE) - await redis.del(hub.CELERY_DEFAULT_QUEUE) - await hub.redisPool.release(redis) - return [hub, closeHub] -} - -describe('queue', () => { - test('process event queue', async () => { - const [hub, closeHub] = await createTestHub() - const redis = await hub.redisPool.acquire() - - // Nothing in the redis queue - const queue1 = await redis.llen(hub.PLUGINS_CELERY_QUEUE) - expect(queue1).toBe(0) - - const kwargs = { - distinct_id: 'hedgehog', - ip: null, - site_url: 'hedgehogs.com', - data: { pineapple: 1 }, - team_id: 1234, - now: 'now', - sent_at: 'later', - } - const args = Object.values(kwargs) - - const client = new Client(hub.db, hub.PLUGINS_CELERY_QUEUE) - for (let i = 0; i < 6; i++) { - client.sendTask('posthog.tasks.process_event.process_event_with_plugins', args, {}) - } - - await delay(1000) - - expect(await redis.llen(hub.PLUGINS_CELERY_QUEUE)).toBe(6) - const fakePiscina = { run: jest.fn() } as any - const queue = (await startQueues(hub, fakePiscina, {})).ingestion - await advanceOneTick() - - await delay(1000) - - expect(await redis.llen(hub.PLUGINS_CELERY_QUEUE)).not.toBe(6) - - await queue!.pause() - - expect(ingestEvent).toHaveBeenCalledTimes(6) - const { data, ...expected_args } = kwargs - expect(ingestEvent).toHaveBeenLastCalledWith( - expect.anything(), - expect.anything(), - expect.objectContaining({ ...expected_args, ...data, uuid: expect.anything() }), - expect.anything() - ) - - await queue!.stop() - await hub.redisPool.release(redis) - await closeHub() - }) - - describe('capabilities', () => { - let hub: Hub - let piscina: Piscina - let closeHub: () => Promise - - beforeEach(async () => { - ;[hub, closeHub] = await createHub({ - LOG_LEVEL: LogLevel.Warn, - }) - piscina = { run: jest.fn() } as any - }) - - afterEach(async () => { - await closeHub() - }) - - it('starts ingestion queue by default', async () => { - const queues = await startQueues(hub, piscina) - - expect(queues).toEqual({ - ingestion: expect.any(KafkaQueue), - }) - }) - - it('handles ingestion being turned off', async () => { - hub.capabilities.ingestion = false - - const queues = await startQueues(hub, piscina) - - expect(queues).toEqual({ - ingestion: null, - }) - }) - }) -}) diff --git a/plugin-server/tests/shared/capabilities.test.ts b/plugin-server/tests/shared/capabilities.test.ts new file mode 100644 index 000000000000..33a31aefeb42 --- /dev/null +++ b/plugin-server/tests/shared/capabilities.test.ts @@ -0,0 +1,45 @@ +import Piscina from '@posthog/piscina' + +import { KafkaQueue } from '../../src/main/ingestion-queues/kafka-queue' +import { startQueues } from '../../src/main/ingestion-queues/queue' +import { Hub, LogLevel } from '../../src/types' +import { createHub } from '../../src/utils/db/hub' + +jest.mock('../../src/main/ingestion-queues/kafka-queue') + +describe('queue', () => { + describe('capabilities', () => { + let hub: Hub + let piscina: Piscina + let closeHub: () => Promise + + beforeEach(async () => { + ;[hub, closeHub] = await createHub({ + LOG_LEVEL: LogLevel.Warn, + }) + piscina = { run: jest.fn() } as any + }) + + afterEach(async () => { + await closeHub() + }) + + it('starts ingestion queue by default', async () => { + const queues = await startQueues(hub, piscina) + + expect(queues).toEqual({ + ingestion: expect.any(KafkaQueue), + }) + }) + + it('handles ingestion being turned off', async () => { + hub.capabilities.ingestion = false + + const queues = await startQueues(hub, piscina) + + expect(queues).toEqual({ + ingestion: null, + }) + }) + }) +})