|
| 1 | +import { validateIamRoleArnFormat, send } from '../utils' |
| 2 | +import { Payload } from '../send/generated-types' |
| 3 | +import { KinesisClient, PutRecordsCommand } from '@aws-sdk/client-kinesis' |
| 4 | +import { assumeRole } from '../../../lib/AWS/sts' |
| 5 | +import { IntegrationError } from '@segment/actions-core' |
| 6 | +import { Logger } from '@segment/actions-core/destination-kit' |
| 7 | + |
| 8 | +describe('validateIamRoleArnFormat', () => { |
| 9 | + it('should return true for a valid IAM Role ARN', () => { |
| 10 | + const validArns = [ |
| 11 | + 'arn:aws:iam::123456789012:role/MyRole', |
| 12 | + 'arn:aws:iam::000000000000:role/service-role/My-Service_Role', |
| 13 | + 'arn:aws:iam::987654321098:role/path/to/MyRole', |
| 14 | + 'arn:aws:iam::111122223333:role/MyRole-With.Special@Chars_+=,.' |
| 15 | + ] |
| 16 | + |
| 17 | + for (const arn of validArns) { |
| 18 | + expect(validateIamRoleArnFormat(arn)).toBe(true) |
| 19 | + } |
| 20 | + }) |
| 21 | + |
| 22 | + it('should return false for an ARN with invalid prefix', () => { |
| 23 | + const invalidArn = 'arn:aws:s3::123456789012:role/MyRole' |
| 24 | + expect(validateIamRoleArnFormat(invalidArn)).toBe(false) |
| 25 | + }) |
| 26 | + |
| 27 | + it('should return false if missing account ID', () => { |
| 28 | + const invalidArn = 'arn:aws:iam:::role/MyRole' |
| 29 | + expect(validateIamRoleArnFormat(invalidArn)).toBe(false) |
| 30 | + }) |
| 31 | + |
| 32 | + it('should return false if account ID is not 12 digits', () => { |
| 33 | + const invalidArns = ['arn:aws:iam::12345:role/MyRole', 'arn:aws:iam::1234567890123:role/MyRole'] |
| 34 | + for (const arn of invalidArns) { |
| 35 | + expect(validateIamRoleArnFormat(arn)).toBe(false) |
| 36 | + } |
| 37 | + }) |
| 38 | + |
| 39 | + it('should return false if missing "role/" segment', () => { |
| 40 | + const invalidArn = 'arn:aws:iam::123456789012:MyRole' |
| 41 | + expect(validateIamRoleArnFormat(invalidArn)).toBe(false) |
| 42 | + }) |
| 43 | + |
| 44 | + it('should return false if role name contains invalid characters', () => { |
| 45 | + const invalidArns = [ |
| 46 | + 'arn:aws:iam::123456789012:role/My Role', // space |
| 47 | + 'arn:aws:iam::123456789012:role/MyRole#InvalidChar' |
| 48 | + ] |
| 49 | + for (const arn of invalidArns) { |
| 50 | + expect(validateIamRoleArnFormat(arn)).toBe(false) |
| 51 | + } |
| 52 | + }) |
| 53 | + |
| 54 | + it('should return false for empty or null values', () => { |
| 55 | + expect(validateIamRoleArnFormat('')).toBe(false) |
| 56 | + // @ts-expect-error testing invalid input type |
| 57 | + expect(validateIamRoleArnFormat(null)).toBe(false) |
| 58 | + // @ts-expect-error testing invalid input type |
| 59 | + expect(validateIamRoleArnFormat(undefined)).toBe(false) |
| 60 | + }) |
| 61 | +}) |
| 62 | + |
| 63 | +jest.mock('@aws-sdk/client-kinesis') |
| 64 | +jest.mock('../../../lib/AWS/sts') |
| 65 | + |
| 66 | +const mockSend = jest.fn() |
| 67 | +const mockLogger: Partial<Logger> = { |
| 68 | + crit: jest.fn(), |
| 69 | + info: jest.fn(), |
| 70 | + warn: jest.fn() |
| 71 | +} |
| 72 | + |
| 73 | +describe('Kinesis send', () => { |
| 74 | + const mockSettings = { |
| 75 | + iamRoleArn: 'arn:aws:iam::123456789012:role/TestRole', |
| 76 | + iamExternalId: 'ext-id' |
| 77 | + } |
| 78 | + |
| 79 | + const mockPayloads: Payload[] = [ |
| 80 | + { |
| 81 | + streamName: 'test-stream', |
| 82 | + awsRegion: 'us-east-1', |
| 83 | + partitionKey: 'pk-1', |
| 84 | + payload: { data: 'test message' }, |
| 85 | + max_batch_size: 500, |
| 86 | + batch_keys: ['awsRegion'] |
| 87 | + } |
| 88 | + ] |
| 89 | + |
| 90 | + beforeEach(() => { |
| 91 | + jest.clearAllMocks() |
| 92 | + ;(assumeRole as jest.Mock).mockResolvedValue({ |
| 93 | + accessKeyId: 'mockAccess', |
| 94 | + secretAccessKey: 'mockSecret', |
| 95 | + sessionToken: 'mockToken' |
| 96 | + }) |
| 97 | + ;(KinesisClient as unknown as jest.Mock).mockImplementation(() => ({ |
| 98 | + send: mockSend |
| 99 | + })) |
| 100 | + }) |
| 101 | + |
| 102 | + it('should throw IntegrationError if partitionKey is missing', async () => { |
| 103 | + const invalidPayload = [ |
| 104 | + { ...mockPayloads[0], partitionKey: '' } // missing partitionKey |
| 105 | + ] |
| 106 | + |
| 107 | + await expect(send(mockSettings, invalidPayload, undefined, mockLogger as Logger)).rejects.toThrow(IntegrationError) |
| 108 | + |
| 109 | + expect(mockLogger.crit).not.toHaveBeenCalled() |
| 110 | + }) |
| 111 | + |
| 112 | + it('should create Kinesis client and send records successfully', async () => { |
| 113 | + mockSend.mockResolvedValueOnce({ Records: [] }) |
| 114 | + |
| 115 | + await send(mockSettings, mockPayloads, undefined, mockLogger as Logger) |
| 116 | + |
| 117 | + expect(assumeRole).toHaveBeenCalledWith( |
| 118 | + mockSettings.iamRoleArn, |
| 119 | + mockSettings.iamExternalId, |
| 120 | + expect.any(String) // region |
| 121 | + ) |
| 122 | + |
| 123 | + expect(KinesisClient).toHaveBeenCalledWith( |
| 124 | + expect.objectContaining({ |
| 125 | + region: 'us-east-1', |
| 126 | + credentials: expect.any(Object) |
| 127 | + }) |
| 128 | + ) |
| 129 | + |
| 130 | + expect(mockSend).toHaveBeenCalledWith(expect.any(PutRecordsCommand)) |
| 131 | + }) |
| 132 | + |
| 133 | + it('should log and rethrow error when Kinesis send fails', async () => { |
| 134 | + const error = new Error('Kinesis failure') |
| 135 | + mockSend.mockRejectedValueOnce(error) |
| 136 | + |
| 137 | + await expect(send(mockSettings, mockPayloads, undefined, mockLogger as Logger)).rejects.toThrow('Kinesis failure') |
| 138 | + |
| 139 | + expect(mockLogger.crit).toHaveBeenCalledWith('Failed to send batch to Kinesis:', error) |
| 140 | + }) |
| 141 | +}) |
0 commit comments