From 8a3e308b70ddfff9c0b940eed7a834c8cd85de17 Mon Sep 17 00:00:00 2001 From: DennisKrone Date: Mon, 5 Feb 2024 16:02:39 +0100 Subject: [PATCH] Fix a bug in mongoreader preventing using existing collections (#167) Fix a bug in mongoreader preventing using existing collections (it didnt like being being cast to a boolean) --- locust_plugins/mongoreader.py | 9 ++++++--- test/test_cloudwatch_plugin.py | 3 ++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/locust_plugins/mongoreader.py b/locust_plugins/mongoreader.py index a2242d2..ede5a88 100644 --- a/locust_plugins/mongoreader.py +++ b/locust_plugins/mongoreader.py @@ -34,9 +34,12 @@ def __init__( """ self.timestamp_field = timestamp_field - self.coll: Collection = ( - coll or MongoClient(env["LOCUST_MONGO"])[env["LOCUST_MONGO_DATABASE"]][env["LOCUST_MONGO_COLLECTION"]] - ) + if not coll: + self.coll: Collection = MongoClient(env["LOCUST_MONGO"])[env["LOCUST_MONGO_DATABASE"]][ + env["LOCUST_MONGO_COLLECTION"] + ] + else: + self.coll = coll self.cursor: Cursor = self.coll.find(filter, sort=[(self.timestamp_field, 1)]) records_in_buffer = self.cursor._refresh() # trigger fetch immediately instead of waiting for the first next() if not records_in_buffer: diff --git a/test/test_cloudwatch_plugin.py b/test/test_cloudwatch_plugin.py index ee9064f..9bbbad1 100644 --- a/test/test_cloudwatch_plugin.py +++ b/test/test_cloudwatch_plugin.py @@ -10,12 +10,12 @@ setup_logging("INFO", None) + class CloudwatchMock: def __init__(self): self.call_count = 0 self.metrics_dict = {} - def put_metric_data(self, Namespace, MetricData): self.call_count += 1 self.metrics_dict[Namespace] = MetricData @@ -23,6 +23,7 @@ def put_metric_data(self, Namespace, MetricData): cw = CloudwatchMock() + def on_locust_init(environment, **_kwargs): CloudwatchAdapter(environment, "MyExampleService", "perf", cw)