Merge pull request #5 from Healy-Hyperspatial/custom-config

Add custom config

Author: jonhealy1

CHANGELOG.md

@@ -15,6 +15,18 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Fixed
 
 
+
+## [v0.2.0]
+
+### Added
+
+- Taurus config options in cli. [#5](https://github.com/Healy-Hyperspatial/stac-api-load-balancing/pull/5)
+
+### Changed
+
+### Fixed
+
+
 ## [v0.1.1]
 
 ### Added
@@ -35,6 +47,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - First release
 
 
-[Unreleased]: <https://github.com/stac-utils/stac-fastapi-elasticsearch/tree/v0.1.1...main>
+[Unreleased]: <https://github.com/stac-utils/stac-fastapi-elasticsearch/tree/v0.2.0...main>
+[v0.2.0]: <https://github.com/stac-utils/stac-fastapi-elasticsearch/tree/v0.1.1...v0.2.0>
 [v0.1.1]: <https://github.com/stac-utils/stac-fastapi-elasticsearch/tree/v0.1.0...v0.1.1>
 [v0.1.0]: <https://github.com/stac-utils/stac-fastapi-elasticsearch/tree/v0.1.0>

README.md

@@ -22,20 +22,29 @@ Usage: stac-api-load-balancing [OPTIONS]
   This tool facilitates data ingestion, Locust load testing, and Taurus
   performance testing against a specified STAC API endpoint.
 
-Args:     
-  ingest (bool): If true, ingests sample data into the specified STAC API.     
-  locust (bool): If true, conducts Locust load testing against the STAC API.     
-  taurus (bool): If true, performs Taurus performance testing against the STAC API.     
-  api_address (str): The URL of the STAC API for testing.  
+  Args:     
+  ingest (bool): If True, ingest sample data into the specified STAC API.     
+  locust (bool): If True, execute Locust load tests against the
+  specified STAC API.     
+  taurus (bool): If True, perform Taurus performance
+  testing with custom settings against the specified STAC API.     
+  concurrency (int): Specifies the number of concurrent users for Taurus testing. Default is 10.      
+  ramp_up (str): Specifies the ramp-up period for Taurus testing, in Taurus notation (e.g., '1m' for 1 minute). Default is '1m'.  
+  iterations (int): Specifies the number of iterations each virtual user will
+  execute in Taurus testing. Default is 100.     
+  api_address (str): The base URL of the STAC API to be tested.
 
 Options:
-  --version               Show the version and exit.
-  -a, --api-address TEXT  Specify the STAC API URL to test against.
-  -t, --taurus            Run the Taurus wrapper for performance testing
-                          against the STAC API.
-  -l, --locust            Run Locust load tests against the STAC API.
-  -i, --ingest            Ingest sample data into the STAC API.
-  --help                  Show this message and exit.
+  -i, --ingest               Ingest sample data into the STAC API.
+  -l, --locust               Run Locust load tests against the STAC API.
+  -t, --taurus               Run the Taurus wrapper for performance testing
+                             against the STAC API.
+  -c, --concurrency INTEGER  Number of concurrent users for Taurus option.
+  -r, --ramp-up TEXT         Ramp-up time for Taurus option.
+  -n, --iterations INTEGER   Number of iterations for Taurus option.
+  -a, --api-address TEXT     Specify the STAC API URL to test against.
+  --version                  Show the version and exit.
+  --help                     Show this message and exit.
 ```
 
 ## Ingest test data - http://localhost:8084 is just an example url

docker-compose.yml

@@ -38,7 +38,6 @@ services:
     ports:
       - "8083:8082"
     volumes:
-      - ./stac_fastapi:/app/stac_fastapi
       - ./scripts:/app/scripts
     depends_on:
       - pgstac

setup.py

@@ -1,7 +1,7 @@
 """stac-api-load-balancing setup.py."""
 from setuptools import setup
 
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
@@ -20,7 +20,6 @@
         "bzt",
     ],
     packages=["stac_api_load_balancing"],
-    # package_data={"stac_api_load_balancing": ["config_files/*", "data_loader/*", "data_loader/setup_data/*"]},
     entry_points={
         "console_scripts": ["stac-api-load-balancing=stac_api_load_balancing.cli:main"]
     },

stac_api_load_balancing/cli.py

@@ -10,12 +10,22 @@
 from .data_loader import data_loader
 
 
-def generate_taurus_config(api_url):
-    """Generate a Taurus configuration file with the specified API URL."""
+def generate_taurus_config(api_url: str, concurrency: int, ramp_up, iterations) -> str:
+    """
+    Generate a custom Taurus configuration file based on the specified settings.
+
+    Args:
+        api_url (str): The base URL for the STAC API to be tested.
+        concurrency (int): The number of concurrent users to simulate.
+        ramp_up (str): The duration over which to ramp up the load test.
+        iterations (int): The total number of iterations to perform.
+
+    Returns:
+        str: The path to the generated Taurus configuration file.
+    """
     template_path = pkg_resources.resource_filename(
         __name__, "config_files/taurus_locust.yml"
     )
-    print("template path: ", template_path)
     locustfile_path = pkg_resources.resource_filename(
         __name__, "config_files/locustfile.py"
     )
@@ -28,6 +38,9 @@ def generate_taurus_config(api_url):
             config = yaml.safe_load(file)
 
         # Setting the script path directly to where `locustfile.py` is expected to be within the package
+        config["execution"][0]["concurrency"] = concurrency
+        config["execution"][0]["ramp-up"] = ramp_up
+        config["execution"][0]["iterations"] = iterations
         config["scenarios"]["default"]["script"] = locustfile_path
         config["scenarios"]["default"]["default-address"] = api_url
 
@@ -40,6 +53,7 @@ def generate_taurus_config(api_url):
         return None
 
 
+@click.command()
 @click.option(
     "-i", "--ingest", is_flag=True, help="Ingest sample data into the STAC API."
 )
@@ -52,50 +66,77 @@ def generate_taurus_config(api_url):
     is_flag=True,
     help="Run the Taurus wrapper for performance testing against the STAC API.",
 )
+@click.option(
+    "-c",
+    "--concurrency",
+    default=10,
+    help="Number of concurrent users for Taurus option.",
+    type=int,
+)
+@click.option("-r", "--ramp-up", default="1m", help="Ramp-up time for Taurus option.")
+@click.option(
+    "-n",
+    "--iterations",
+    default=100,
+    help="Number of iterations for Taurus option.",
+    type=int,
+)  # Changed the flag to -n to avoid conflict
 @click.option(
     "-a",
     "--api-address",
     default="http://localhost:8080",
     help="Specify the STAC API URL to test against.",
 )
-@click.command()
-@click.version_option(version="0.1.0")
-def main(ingest, locust, taurus, api_address):
+@click.version_option(version="0.2.0")
+def main(
+    ingest: bool,
+    locust: bool,
+    taurus: bool,
+    api_address: str,
+    concurrency: int,
+    ramp_up: str,
+    iterations: int,
+):
     """
     Entry point for the stac-api-load-balancing CLI tool.
 
     This tool facilitates data ingestion, Locust load testing, and Taurus performance testing
     against a specified STAC API endpoint.
 
     Args:
-        ingest (bool): If true, ingests sample data into the specified STAC API.
-        locust (bool): If true, conducts Locust load testing against the STAC API.
-        taurus (bool): If true, performs Taurus performance testing against the STAC API.
-        api_address (str): The URL of the STAC API for testing.
+        ingest (bool): If True, ingest sample data into the specified STAC API.
+        locust (bool): If True, execute Locust load tests against the specified STAC API.
+        taurus (bool): If True, perform Taurus performance testing with custom settings against the specified STAC API.
+        concurrency (int): Specifies the number of concurrent users for Taurus testing. Default is 10.
+        ramp_up (str): Specifies the ramp-up period for Taurus testing, in Taurus notation (e.g., '1m' for 1 minute). Default is '1m'.
+        iterations (int): Specifies the number of iterations each virtual user will execute in Taurus testing. Default is 100.
+        api_address (str): The base URL of the STAC API to be tested.
     """
     os.environ["LOCUST_HOST"] = api_address
 
     if ingest:
+        # Load data into the STAC API
         data_loader.load_items(stac_api_base_url=api_address)
     elif locust:
+        # Execute Locust load tests
+        locust_file_path = pkg_resources.resource_filename(
+            __name__, "config_files/locustfile.py"
+        )
         subprocess.run(
-            [
-                "locust",
-                "--locustfile",
-                pkg_resources.resource_filename(__name__, "config_files/locustfile.py"),
-                "--host",
-                api_address,
-            ],
+            ["locust", "--locustfile", locust_file_path, "--host", api_address],
             check=True,
         )
     elif taurus:
-        config_file_path = generate_taurus_config(api_address)
+        # Generate and run a custom Taurus configuration for performance testing
+        config_file_path = generate_taurus_config(
+            api_address, concurrency, ramp_up, iterations
+        )
         if config_file_path:
             try:
                 subprocess.run(["bzt", config_file_path], check=True)
             finally:
                 if os.path.exists(config_file_path):
-                    os.remove(config_file_path)  # Cleanup the temporary config file
+                    os.remove(config_file_path)  # Cleanup after running
 
 
 if __name__ == "__main__":