locust-test.py

import psycopg2
from locust import User, between, task, wait_time

class MyUser(User):
    wait_time = between(5, 15)
    def init(self, *args, **kwargs):
        super().init(*args, **kwargs)
        self.conn = psycopg2.connect(
        host="localhost",
        database="mydatabase",
        user="myuser",
        password="mypassword"
        )
    def on_stop(self):
        self.conn.close()

    @task
    def select_query(self):
        with self.conn.cursor() as cur:
            cur.execute("SELECT * FROM mytable")
            rows = cur.fetchall()
            assert len(rows) > 0
    @task
    def insert_query(self):
        with self.conn.cursor() as cur:
            cur.execute("INSERT INTO mytable (col1, col2, col3) VALUES ('value1', 'value2', 'value3')")
            self.conn.commit()