clean up locustfile

examples/postgres/locustfile.py

@@ -5,31 +5,24 @@
 import random
 
 
-class UserTasks(TaskSet):
+class PostgresLocust(PostgresUser):
     @task
     def run_select_query(self):
         self.client.execute_query(
             "SELECT * FROM loadtesting.invoice WHERE amount > 500",
         )
 
-    @task(3)
+    @task
     def run_update_query(self):
         random_amount = random.randint(1, 12)
         self.client.execute_query(
             f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10",
         )
 
-
-class PostgresLocust(PostgresUser):
-    tasks = [UserTasks]
-    min_wait = 0
-    max_wait = 3
-    wait_time = between(min_wait, max_wait)
-
     # Use environment variables or default values
     PGHOST = os.getenv("PGHOST", "localhost")
     PGPORT = os.getenv("PGPORT", "5432")
-    PGDATABASE = os.getenv("PGDATABASE", "postgres")
+    PGDATABASE = os.getenv("PGDATABASE", "test_db")
     PGUSER = os.getenv("PGUSER", "postgres")
     PGPASSWORD = os.getenv("PGPASSWORD", "postgres")
 

locust/contrib/postgres.py

@@ -1,24 +1,18 @@
-from locust import TaskSet, User, events
+from locust import User, events
 
 import time
 
 import psycopg
 
 
-def create_conn(conn_string):
-    return psycopg.connect(conn_string)
-
-
 class PostgresClient:
     def __init__(self, conn_string):
-        self.conn_string = conn_string
-        self.connection = create_conn(conn_string)
+        self.connection = psycopg.connect(conn_string)
 
     def execute_query(self, query):
         start_time = time.time()
         try:
-            cursor = self.connection.cursor()
-            cursor.execute(query)
+            self.connection.execute(query)
             response_time = int((time.time() - start_time) * 1000)
             events.request.fire(
                 request_type="postgres_success",
@@ -35,7 +29,6 @@ def execute_query(self, query):
                 response_length=0,
                 exception=e,
             )
-            print(f"error: {e}")
 
     def close(self):
         self.connection.close()