|
18 | 18 | */ |
19 | 19 |
|
20 | 20 | const Pulsar = require('../index'); |
| 21 | +const httpUtils = require('./http_utils'); |
| 22 | + |
| 23 | +function getPartition(msgId) { |
| 24 | + // The message id string is in the format of "entryId,ledgerId,partition,batchIndex" |
| 25 | + return Number(msgId.toString().split(',')[2]); |
| 26 | +} |
21 | 27 |
|
22 | 28 | (() => { |
23 | 29 | describe('Producer', () => { |
@@ -156,5 +162,82 @@ const Pulsar = require('../index'); |
156 | 162 | await producer2.close(); |
157 | 163 | }); |
158 | 164 | }); |
| 165 | + describe('Message Routing', () => { |
| 166 | + test('Custom Message Router', async () => { |
| 167 | + const topic = `test-custom-router-${Date.now()}`; |
| 168 | + const numPartitions = 3; |
| 169 | + const response = await httpUtils.createPartitionedTopic(topic, numPartitions); |
| 170 | + expect(response.statusCode).toBe(204); |
| 171 | + |
| 172 | + const producer = await client.createProducer({ |
| 173 | + topic, |
| 174 | + batchingMaxMessages: 2, |
| 175 | + messageRouter: (message, topicMetadata) => parseInt(message.getPartitionKey(), 10) |
| 176 | + % topicMetadata.numPartitions, |
| 177 | + messageRoutingMode: 'CustomPartition', |
| 178 | + }); |
| 179 | + |
| 180 | + const promises = []; |
| 181 | + const numMessages = 5; |
| 182 | + for (let i = 0; i < numMessages; i += 1) { |
| 183 | + const sendPromise = producer.send({ |
| 184 | + partitionKey: `${i}`, |
| 185 | + data: Buffer.from(`msg-${i}`), |
| 186 | + }); |
| 187 | + await sendPromise; |
| 188 | + promises.push(sendPromise); |
| 189 | + } |
| 190 | + try { |
| 191 | + const allMsgIds = await Promise.all(promises); |
| 192 | + console.log(`All messages have been sent. IDs: ${allMsgIds.join(', ')}`); |
| 193 | + for (let i = 0; i < allMsgIds.length; i += 1) { |
| 194 | + // The message id string is in the format of "entryId,ledgerId,partition,batchIndex" |
| 195 | + const partition = getPartition(allMsgIds[i]); |
| 196 | + expect(i % numPartitions).toBe(partition); |
| 197 | + } |
| 198 | + } catch (error) { |
| 199 | + console.error('One or more messages failed to send:', error); |
| 200 | + } |
| 201 | + }, 30000); |
| 202 | + test('Exception in router', async () => { |
| 203 | + const topic = `test-exception-in-router-${Date.now()}`; |
| 204 | + const numPartitions = 2; |
| 205 | + const response = await httpUtils.createPartitionedTopic(topic, numPartitions); |
| 206 | + expect(response.statusCode).toBe(204); |
| 207 | + const producer = await client.createProducer({ |
| 208 | + topic, |
| 209 | + messageRouter: (message, topicMetadata) => { |
| 210 | + throw new Error('Custom error in message router'); |
| 211 | + }, |
| 212 | + messageRoutingMode: 'CustomPartition', |
| 213 | + }); |
| 214 | + await expect( |
| 215 | + producer.send({ data: Buffer.from('test') }), |
| 216 | + ).rejects.toThrow('Failed to send message: UnknownError'); |
| 217 | + }, 30000); |
| 218 | + test('Not CustomPartition', async () => { |
| 219 | + const topic = `test-not-custom-part-${Date.now()}`; |
| 220 | + const numPartitions = 2; |
| 221 | + const response = await httpUtils.createPartitionedTopic(topic, numPartitions); |
| 222 | + expect(response.statusCode).toBe(204); |
| 223 | + |
| 224 | + let index = 0; |
| 225 | + const producer = await client.createProducer({ |
| 226 | + topic, |
| 227 | + messageRouter: (_, topicMetadata) => { |
| 228 | + const result = index % topicMetadata.numPartitions; |
| 229 | + index += 1; |
| 230 | + return result; |
| 231 | + }, |
| 232 | + messageRoutingMode: 'UseSinglePartition', |
| 233 | + }); |
| 234 | + const partitions = new Set(); |
| 235 | + for (let i = 0; i < 10; i += 1) { |
| 236 | + const msgId = await producer.send({ data: Buffer.from('msg') }); |
| 237 | + partitions.add(getPartition(msgId)); |
| 238 | + } |
| 239 | + expect(partitions.size).toBe(1); |
| 240 | + }, 30000); |
| 241 | + }); |
159 | 242 | }); |
160 | 243 | })(); |
0 commit comments