mongodb_locust.py

from locust import User, task, constant, tag
from pymongo import MongoClient
from time import perf_counter


class Mongo:
    conn = ""

    def __init__(self):
        cs = 'mongodb://mongodb-0:27017,mongodb-1:27011,mongodb-2:27012,mongodb-3:27013,mongodb-arbiter-0:27014'
        self.conn = MongoClient(cs, connect=False, replicaSet='rs0')

    def get_conn(self):
        return self.conn

mongo_client = Mongo()

class Mongouser(User):
    wait_time = constant(1)

    @task
    @tag("GET")
    def aggregate_test(self):
        with mongo_client.get_conn() as cx:
            db = cx.test

            start_time = perf_counter()

            # dummy aggregate example
            pipeline = [
                {"$unwind": "$tags"},
                {"$group": {"_id": "$tags", "count": {"$sum": 1}}},
                {"$sort": SON([("count", -1), ("_id", -1)])}
            ]
            data = list(db.test_collection.aggregate(pipeline))
            duration = perf_counter() - start_time

            if data:
                self.environment.events.request_success.fire(request_type="GET", name="Aggregate test",
                                                             response_time=duration, response_length=len(data))
            else:
                self.environment.events.request_failure.fire(request_type="GET", name="Aggregate test",
                                                             response_time=duration, response_length=0)

    @task
    @tag("GET")
    def fetch_test(self):
        with mongo_client.get_conn() as cx:
            db = cx.test
            start_time = perf_counter()
            try:
                doc_count = db.test_collection.count_documents({})
                duration = perf_counter() - start_time

                self.environment.events.request_success.fire(request_type="GET", name="Fetch documents",
                                                             response_time=duration, response_length=doc_count)
            except Exception as e:
                duration = perf_counter() - start_time
                print(e)
                self.environment.events.request_failure.fire(request_type="GET", name="Fetch documents",
                                                             response_time=duration, response_length=0, exception=e)