benchmark for AWS Kinesis Data Stream

locust-kinesis.py

# -*- coding: utf-8 -*-

import time
import datetime
import boto3
from locust import HttpLocust, TaskSet, task
from locust.events import request_success


class UserBehavior(TaskSet):
    
    @task
    def index(self):
        self.client.get('/')

class Kinesis(TaskSet):

    def on_start(self):
        self.kinesis_client = boto3.client('kinesis', region_name='us-west-2')

    
    @task
    def put(self):
        stream_name = 'my-stream'
        
        start_at = time.time()
        response = self.kinesis_client.put_record(
            StreamName=stream_name,
            Data='hogehoge\n',
            PartitionKey='1'
        )
        request_success.fire(
            request_type='PUT',
            name=stream_name,
            response_time=int((time.time() - start_at) * 1000),
            response_length=len(response)
        )


class WebsiteUser(HttpLocust):
    # task_set = UserBehavior
    task_set = Kinesis
    min_wait = 0
    max_wait = 10