diff --git a/examples/postgres/locustfile.py b/examples/postgres/locustfile.py index 6904ec88b0..fca4f9eee9 100644 --- a/examples/postgres/locustfile.py +++ b/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") diff --git a/locust/contrib/postgres.py b/locust/contrib/postgres.py index 7444666dd0..7cd1ead938 100644 --- a/locust/contrib/postgres.py +++ b/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()