Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

test(plugin-server): Improve the robustness of overflow handling tests #21315

Merged
merged 3 commits into from
Apr 4, 2024
Merged

test(plugin-server): Improve the robustness of overflow handling tests #21315

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dc36a05
Pass `produce` calls through to through to real producer.
tkaemming Apr 3, 2024
ceb218a
Better align tests with actual interfaces/behavior
tkaemming Apr 3, 2024
f13163e
Just use `computeKey` which is already defined
tkaemming Apr 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 47 additions & 40 deletions plugin-server/tests/main/ingestion-queues/analytics-events-ingestion-consumer.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
import { Message } from 'node-rdkafka'

import { buildStringMatcher } from '../../../src/config/config'
import { KAFKA_EVENTS_PLUGIN_INGESTION, KAFKA_EVENTS_PLUGIN_INGESTION_OVERFLOW } from '../../../src/config/kafka-topics'
import {
eachBatchParallelIngestion,
IngestionOverflowMode,
} from '../../../src/main/ingestion-queues/batch-processing/each-batch-ingestion'
import { Hub } from '../../../src/types'
import { createHub } from '../../../src/utils/db/hub'
import { ConfiguredLimiter } from '../../../src/utils/token-bucket'
import { captureIngestionWarning } from './../../../src/worker/ingestion/utils'

Expand Down Expand Up @@ -47,58 +51,59 @@ const captureEndpointEvent2 = {
}

describe('eachBatchParallelIngestion with overflow reroute', () => {
let hub: Hub
let closeServer: () => Promise<void>
let queue: any

function createBatchWithMultipleEventsWithKeys(events: any[], timestamp?: any): any {
return events.map((event) => ({
function makeMessageKey(event: any): string {
return event.token + ':' + event.distinct_id
}

function createBatchWithMultipleEvents(events: any[], timestamp?: any, withKey: boolean = true): Message[] {
return events.map((event, i) => ({
partition: 0,
topic: KAFKA_EVENTS_PLUGIN_INGESTION,
value: JSON.stringify(event),
value: Buffer.from(JSON.stringify(event)),
timestamp,
offset: event.offset,
key: event.team_id + ':' + event.distinct_id,
offset: i,
key: withKey ? makeMessageKey(event) : null,
size: 0, // irrelevant, but needed for type checking
}))
}

beforeEach(() => {
beforeEach(async () => {
;[hub, closeServer] = await createHub()
queue = {
bufferSleep: jest.fn(),
pluginsServer: {
INGESTION_CONCURRENCY: 4,
kafkaProducer: {
produce: jest.fn(),
},
db: 'database',
},
pluginsServer: hub,
}
jest.mock('./../../../src/worker/ingestion/event-pipeline/runner')
})

afterEach(async () => {
await closeServer()
jest.clearAllMocks()
})

it('reroutes events with no key to OVERFLOW topic', async () => {
const batch = [
{
partition: 0,
topic: KAFKA_EVENTS_PLUGIN_INGESTION,
value: JSON.stringify(captureEndpointEvent1),
timestamp: captureEndpointEvent1['timestamp'],
offset: captureEndpointEvent1['offset'],
key: null,
token: 'ok',
},
]
const now = Date.now()
const [message] = createBatchWithMultipleEvents(
[captureEndpointEvent1],
now,
false // act as if this message was intended to be routed to overflow by capture
)

const consume = jest.spyOn(ConfiguredLimiter, 'consume').mockImplementation(() => false)
const produce = jest.spyOn(queue.pluginsServer.kafkaProducer, 'produce')

const tokenBlockList = buildStringMatcher('another_token,more_token', false)
await eachBatchParallelIngestion(tokenBlockList, batch, queue, IngestionOverflowMode.Reroute)
await eachBatchParallelIngestion(tokenBlockList, [message], queue, IngestionOverflowMode.Reroute)

expect(consume).not.toHaveBeenCalled()
expect(captureIngestionWarning).not.toHaveBeenCalled()
expect(queue.pluginsServer.kafkaProducer.produce).toHaveBeenCalledWith({
expect(produce).toHaveBeenCalledWith({
topic: KAFKA_EVENTS_PLUGIN_INGESTION_OVERFLOW,
value: JSON.stringify(captureEndpointEvent1),
timestamp: captureEndpointEvent1['timestamp'],
offset: captureEndpointEvent1['offset'],
Comment on lines -100 to -101
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These were undefined so they weren't being considered in this assertion.

timestamp can be provided to produce, if we want to preserve it across rerouting.

value: message.value,
key: null,
waitForAck: true,
})
Expand All @@ -109,23 +114,23 @@ describe('eachBatchParallelIngestion with overflow reroute', () => {

it('reroutes excess events to OVERFLOW topic', async () => {
Copy link
Contributor Author

@tkaemming tkaemming Apr 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't quite test this by cherry-picking the problematic change (because it also changed these tests), but #20945 (comment) would have caused this test to hang indefinitely now that it actually produces. This would have let us know something weird was going on earlier.

const now = Date.now()
const batch = createBatchWithMultipleEventsWithKeys([captureEndpointEvent1], now)
const event = captureEndpointEvent1
const [message] = createBatchWithMultipleEvents([event], now)
const consume = jest.spyOn(ConfiguredLimiter, 'consume').mockImplementation(() => false)
const produce = jest.spyOn(queue.pluginsServer.kafkaProducer, 'produce')

const tokenBlockList = buildStringMatcher('another_token,more_token', false)
await eachBatchParallelIngestion(tokenBlockList, batch, queue, IngestionOverflowMode.Reroute)
await eachBatchParallelIngestion(tokenBlockList, [message], queue, IngestionOverflowMode.Reroute)

expect(consume).toHaveBeenCalledWith(
captureEndpointEvent1['token'] + ':' + captureEndpointEvent1['distinct_id'],
makeMessageKey(event), // NOTE: can't use ``message.key`` here as it will already have been mutated
1,
now
)
expect(captureIngestionWarning).not.toHaveBeenCalled()
expect(queue.pluginsServer.kafkaProducer.produce).toHaveBeenCalledWith({
expect(produce).toHaveBeenCalledWith({
topic: KAFKA_EVENTS_PLUGIN_INGESTION_OVERFLOW,
value: JSON.stringify(captureEndpointEvent1),
timestamp: captureEndpointEvent1['timestamp'],
offset: captureEndpointEvent1['offset'],
value: message.value,
key: null,
waitForAck: true,
})
Expand All @@ -136,8 +141,9 @@ describe('eachBatchParallelIngestion with overflow reroute', () => {

it('does not reroute if not over capacity limit', async () => {
const now = Date.now()
const batch = createBatchWithMultipleEventsWithKeys([captureEndpointEvent1, captureEndpointEvent2], now)
const batch = createBatchWithMultipleEvents([captureEndpointEvent1, captureEndpointEvent2], now)
const consume = jest.spyOn(ConfiguredLimiter, 'consume').mockImplementation(() => true)
const produce = jest.spyOn(queue.pluginsServer.kafkaProducer, 'produce')

const tokenBlockList = buildStringMatcher('another_token,more_token', false)
await eachBatchParallelIngestion(tokenBlockList, batch, queue, IngestionOverflowMode.Reroute)
Expand All @@ -153,17 +159,18 @@ describe('eachBatchParallelIngestion with overflow reroute', () => {
now
)
expect(captureIngestionWarning).not.toHaveBeenCalled()
expect(queue.pluginsServer.kafkaProducer.produce).not.toHaveBeenCalled()
expect(produce).not.toHaveBeenCalled()
// Event is processed
expect(runEventPipeline).toHaveBeenCalledTimes(2)
})

it('does drop events from blocked tokens', async () => {
const now = Date.now()
const batch = createBatchWithMultipleEventsWithKeys(
const batch = createBatchWithMultipleEvents(
[captureEndpointEvent1, captureEndpointEvent2, captureEndpointEvent1],
now
)
const produce = jest.spyOn(queue.pluginsServer.kafkaProducer, 'produce')
const consume = jest.spyOn(ConfiguredLimiter, 'consume').mockImplementation(() => true)

const tokenBlockList = buildStringMatcher('mytoken,another_token', false)
Expand All @@ -176,7 +183,7 @@ describe('eachBatchParallelIngestion with overflow reroute', () => {
now
)
expect(captureIngestionWarning).not.toHaveBeenCalled()
expect(queue.pluginsServer.kafkaProducer.produce).not.toHaveBeenCalled()
expect(produce).not.toHaveBeenCalled()
expect(runEventPipeline).toHaveBeenCalledTimes(1)
})
})
Loading