Merge pull request #20 from pjcalvo/release-1.0

fixed succes error + renamed tags to extra tags + make db creation att…

Author: pjcalvo

.github/workflows/python-publish.yml

@@ -0,0 +1,39 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/python-test.yml

@@ -0,0 +1,33 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Python package
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Test with pytest
+      run: |
+        python test_suite.py

README.md

@@ -32,15 +32,16 @@ def on_locust_init(environment, **_kwargs):
     """
     # this settings matches the given docker-compose file
     influxDBSettings = InfluxDBSettings(
-        influx_host = 'localhost',
-        influx_port = '8086',
+        host = 'localhost',
+        port = '8086',
         user = 'admin',
         pwd = 'pass',
         database = 'test-project'
 
         # optional global tags to be added to each metric sent to influxdb
-        tags = {
-            'some-global-tag': 'some-global-tag-value',
+        additional_tags = {
+            'environment': 'test',
+            'some_other_tag': 'tag_value',
         }
     )
     # start listerner with the given configuration

example/locust.conf

@@ -1,6 +1,6 @@
 locustfile = locustfiles/locustfile.py # file to run
 headless = true # to run with the locust UI
-host = https://hoodoo.digital # base url for the domains
+host = https://pjcalvo.github.io # base url for the domains
 users = 5 # amout of users that will run
 master = false # to set master/worker state
 spawn-rate = 1 # spawn rate per second

example/locustfiles/locustfile.py

@@ -10,15 +10,16 @@ def on_locust_init(environment, **_kwargs):
     """
     # this settings matches the given docker-compose file
     influxDBSettings = InfluxDBSettings(
-        influx_host = 'localhost',
-        influx_port = '8086',
+        host = 'localhost',
+        port = '8086',
         user = 'admin',
         pwd = 'pass',
         database = 'test-project',
 
-        # optional global tags to be added to each metric sent to influxdb
-        tags = {
-            'some-global-tag': 'some-global-tag-value',
+        # additional_tags tags to be added to each metric sent to influxdb
+        additional_tags = {
+            'environment': 'test',
+            'some_other_tag': 'tag_value',
         }
     )
     # start listerner with the given configuration
@@ -35,14 +36,15 @@ def home_page(self):
             if response.status_code != 200:
                 response.failure("Got wrong response")
 
-    @tag('connectors')
+    # purposely fail finding the text
+    @tag('about_page')
     @task(1)
-    def workfront_connector(self):
-        with self.client.get("/connectors/workfront", catch_response=True) as response:
-            if 'Do More Work, Faster' not in response.text:
-                response.failure("Expected test was not displayed")
+    def about_page(self):
+        with self.client.get("/about", catch_response=True) as response:
+            if 'native nicaraguan' not in response.text:
+                response.failure("Expected text was not displayed")
 
-    
+    # method that runs once per spawned user. Useful for setting up test data or attributes
     def on_start(self):
         print('New user was spawned')
 

example/requirements.txt

@@ -1,3 +1,3 @@
-locust_influxdb_listener == 0.0.9
+locust_influxdb_listener == 1.0.1
 influxdb==5.3.1
 locust==1.5.0

locust_influxdb_listener/__init__.py

@@ -1,3 +1,5 @@
+from typing import Optional, Dict
+
 import atexit
 import gevent
 import logging
@@ -6,40 +8,54 @@
 from datetime import datetime
 
 from influxdb import InfluxDBClient
-from locust.exception import InterruptTaskSet
 from requests.exceptions import HTTPError
 import locust.env
+from urllib3 import HTTPConnectionPool
 
-log = logging.getLogger('locust_influx')
+log = logging.getLogger('locust_influxdb_listener')
 
 
 class InfluxDBSettings:
     """
-    Store influxdb settings
+    Store InfluxDB settings for a data connection.
     """
+
     def __init__(
         self, 
-        influx_host: str = 'localhost', 
-        influx_port: int = 8086, 
+        host: str = 'localhost', 
+        port: int = 8086, 
         user: str = 'admin', 
         pwd: str = 'pass',
         database: str = 'default',
         interval_ms: int = 1000,
         ssl: bool = False,
         verify_ssl: bool = False,
-        create_database: bool = False,
-        tags: dict = {}
+        additional_tags: dict = {},
+        influx_host: Optional[str] = 'localhost', 
+        influx_port: Optional[int] = 8086, 
     ):
-        self.influx_host = influx_host
-        self.influx_port = influx_port
+        """
+        Initialize the InfluxDBSettings object with provided or default settings.
+
+        :param host: InfluxDB host address or hostname.
+        :param port: InfluxDB HTTP API port.
+        :param user: InfluxDB username for authentication.
+        :param pwd: InfluxDB password for authentication.
+        :param database: InfluxDB database name for storing data.
+        :param interval_ms: Data sending interval in milliseconds.
+        :param ssl: Enable SSL/TLS for secure data transmission.
+        :param verify_ssl: Verify SSL certificates (only if SSL is enabled).
+        :param additional_tags: Additional tags to include in globally for all data points.
+        """
+        self.host = host if host else influx_host  # Renamed from influx_host
+        self.port = port if port else influx_port  # Renamed from influx_port
         self.user = user
         self.pwd = pwd
         self.database = database
         self.interval_ms = interval_ms
         self.ssl = ssl
         self.verify_ssl = verify_ssl
-        self.create_database = create_database
-        self.tags = tags
+        self.additional_tags = additional_tags
 
 
 class InfluxDBListener: 
@@ -53,28 +69,35 @@ def __init__(
         env: locust.env.Environment,
         influxDbSettings: InfluxDBSettings
     ):
+        """
+        Initialize the InfluxDBListener with the provided Locust environment and InfluxDB settings.
 
-        # flush related attributes
+        :param env: The Locust environment to listen for events in.
+        :param influxDbSettings: Settings for the InfluxDB connection.
+        """
+         
         self.env = env
         self.cache = []
         self.stop_flag = False
         self.interval_ms = influxDbSettings.interval_ms
-        self.tags = influxDbSettings.tags
+        self.additional_tags = influxDbSettings.additional_tags
         # influxdb settings 
         try:
+            # try to connect create the database and switch to it
             self.influxdb_client = InfluxDBClient(
-                host=influxDbSettings.influx_host,
-                port=influxDbSettings.influx_port,
+                host=influxDbSettings.host,
+                port=influxDbSettings.port,
                 username=influxDbSettings.user,
                 password=influxDbSettings.pwd,
-                database=influxDbSettings.database,
                 ssl=influxDbSettings.ssl,
                 verify_ssl=influxDbSettings.verify_ssl
             )
-            if influxDbSettings.create_database:
-                self.influxdb_client.create_database(influxDbSettings.database)
-        except:
-            logging.error('Could not connect to influxdb.')
+            # database is mandatory so we should always try to create it
+            self.influxdb_client.create_database(influxDbSettings.database)
+            self.influxdb_client.switch_database(influxDbSettings.database)
+            
+        except Exception as ex:
+            logging.error(f'Unexpected error: {ex}')
             return
 
         # determine if worker or master
@@ -85,7 +108,7 @@ def __init__(
             # TODO: Get real ID of slaves form locust somehow
             self.node_id = 'worker'
 
-        # start background event to push data to influx
+        # start background event to push data to influxdb
         self.flush_worker = gevent.spawn(self.__flush_cached_points_worker)
         self.test_start(0)
 
@@ -133,14 +156,13 @@ def __register_event(self, node_id: str, user_count: int, event: str, **_kwargs)
         """
 
         time = datetime.utcnow()
-        tags = self.tags
         fields = {
             'node_id': node_id,
             'event': event,
             'user_count': user_count
         }
 
-        point = self.__make_data_point('locust_events', tags, fields, time)
+        point = self.__make_data_point('locust_events', fields, time)
         self.cache.append(point)
 
     def __listen_for_requests_events(self, node_id, measurement, request_type, name,
@@ -154,9 +176,8 @@ def __listen_for_requests_events(self, node_id, measurement, request_type, name,
 
         time = datetime.utcnow()
         was_successful = True
-        if response:
-            # override with response code
-            was_successful = 199 < response.status_code < 400
+        if response is not None:
+            was_successful = (199 < response.status_code < 400) and exception is None            
         tags = {
             'node_id': node_id,
             'request_type': request_type,
@@ -176,7 +197,7 @@ def __listen_for_requests_events(self, node_id, measurement, request_type, name,
             'counter': self.env.stats.num_requests,  # TODO: Review the need of this field
             'user_count': self.env.runner.user_count
         }
-        point = self.__make_data_point(measurement, tags, fields, time)
+        point = self.__make_data_point(measurement, fields, time, tags=tags)
         self.cache.append(point)
 
     def __listen_for_locust_errors(self, node_id, user_instance, exception: Exception = None, tb=None) -> None:
@@ -194,7 +215,7 @@ def __listen_for_locust_errors(self, node_id, user_instance, exception: Exceptio
             'exception': repr(exception),
             'traceback': "".join(traceback.format_tb(tb)),
         }
-        point = self.__make_data_point('locust_exceptions', tags, fields, time)
+        point = self.__make_data_point('locust_exceptions', fields, time, tags=tags)
         self.cache.append(point)
 
 
@@ -210,15 +231,15 @@ def __flush_cached_points_worker(self) -> None:
             self.__flush_points(self.influxdb_client)
             gevent.sleep(self.interval_ms / 1000)
 
-    def __make_data_point(self, measurement: str, tags: dict, fields: dict, time: datetime) -> dict:
+    def __make_data_point(self, measurement: str, fields: dict, time: datetime, tags: Optional[Dict[str,str]] = {}) -> dict:
         """
         Create a list with a single point to be saved to influxdb.
         :param measurement: The measurement where to save this point.
-        :param tags: Dictionary of tags to be saved in the measurement.
         :param fields: Dictionary of field to be saved to measurement.
         :param time: The time os this point.
+        :param tags: Dictionary of tags to be saved in the measurement default to None.
         """
-        return {"measurement": measurement, "tags": {**tags, **self.tags}, "time": time, "fields": fields}
+        return {"measurement": measurement, "tags": {**tags, **self.additional_tags}, "time": time, "fields": fields}
 
 
     def last_flush_on_quitting(self):

requirements.txt

@@ -0,0 +1,2 @@
+influxdb==5.3.1
+locust==1.5.0

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="locust_influxdb_listener", # Replace with your own username
-    version="0.0.9",
+    version="1.0.1",
     author="Pablo Calvo",
     author_email="[email protected]",
     description="Locust.io 1.X influxdb listener",

test_suite.py

@@ -0,0 +1,17 @@
+import unittest
+
+def create_test_suite():
+    loader = unittest.TestLoader()
+    start_dir = "tests" 
+    suite = loader.discover(start_dir, pattern="test_*.py")
+    return suite
+
+if __name__ == "__main__":
+    test_suite = create_test_suite()
+    runner = unittest.TextTestRunner(verbosity=2)
+    result = runner.run(test_suite)
+    
+    if result.wasSuccessful():
+        exit(0)
+    else:
+        exit(1)