From 8633a1c71d3f923fc21a655b9cad06bd60343517 Mon Sep 17 00:00:00 2001
From: Tarun Menon <tarunmenon95@gmail.com>
Date: Thu, 11 Jul 2024 16:19:15 +1000
Subject: [PATCH] add locust slave config

---
 .github/workflows/locust.yaml | 53 ++++++++++++++++++++++++++---------
 locust/docker-entrypoint.sh   | 38 ++++++++++++++++++-------
 2 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/locust.yaml b/.github/workflows/locust.yaml
index d955b48..7e21ec6 100644
--- a/.github/workflows/locust.yaml
+++ b/.github/workflows/locust.yaml
@@ -2,7 +2,9 @@ name: locust
 
 on:
   push:
-    branches: [ master ]
+    branches:
+      - master
+      - feature/*
     paths:
     - 'locust/**'
 
@@ -14,18 +16,43 @@ jobs:
     runs-on: ubuntu-latest
     if: github.ref == 'refs/heads/master' # Running this job only for master branch
 
-    steps:
-    - uses: actions/checkout@v2
+    # steps:
+    # - uses: actions/checkout@v2
 
-    - name: Build image
-      run: docker build --tag $IMAGE_NAME --label "runnumber=${GITHUB_RUN_ID}" locust/
+    # - name: Build image
+    #   run: docker build --tag $IMAGE_NAME --label "runnumber=${GITHUB_RUN_ID}" locust/
 
-    - name: Log into registry
-      run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+    # - name: Log into registry
+    #   run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
 
-    - name: Push image
-      run: |
-        IMAGE_ID=ghcr.io/${{ github.repository_owner }}/$IMAGE_NAME
-        IMAGE_ID=$(echo $IMAGE_ID | tr '[A-Z]' '[a-z]')
-        docker tag $IMAGE_NAME $IMAGE_ID:latest
-        docker push $IMAGE_ID:latest
+    # - name: Push image
+    #   run: |
+    #     IMAGE_ID=ghcr.io/${{ github.repository_owner }}/$IMAGE_NAME
+    #     IMAGE_ID=$(echo $IMAGE_ID | tr '[A-Z]' '[a-z]')
+    #     docker tag $IMAGE_NAME $IMAGE_ID:latest
+    #     docker push $IMAGE_ID:latest
+    
+    steps:
+      - name: Check out the repo
+        uses: actions/checkout@v3
+  
+      - name: Inject slug/short variables
+        uses: rlespinasse/github-slug-action@v3.x  
+  
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+  
+      - name: Login to  GitHub Container Repository
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+  
+      - name: Build and push Container Image to GitHub Container Repository
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ./Dockerfile
+          push: true
+          tags: ghcr.io/base2services/locust:snapshot_${{env.GITHUB_REF_SLUG}}
\ No newline at end of file
diff --git a/locust/docker-entrypoint.sh b/locust/docker-entrypoint.sh
index 56792cf..3a44100 100755
--- a/locust/docker-entrypoint.sh
+++ b/locust/docker-entrypoint.sh
@@ -1,17 +1,35 @@
 #!/bin/sh
 
-set -e
+set -ex
 
-if [ "x" != "${LOCUST_S3_PATH}x" ] ; then
-    aws s3 cp s3://${LOCUST_S3_PATH}/${LOCUST_LOCUSTFILE} . --region ${AWS_REGION}
-    chmod +x ${LOCUST_LOCUSTFILE}
+aws s3 cp s3://${LOCUST_S3_PATH} /locust/ --recursive --region ${AWS_REGION}
+
+LOCUST_MODE=${LOCUST_MODE:-standalone}
+LOCUST_MASTER_BIND_PORT=${LOCUST_MASTER_BIND_PORT:-5557}
+LOCUST_FILE=${LOCUST_FILE:-locustfile.py}
+LOCUST_LOG_LEVEL=${LOCUST_LOG_LEVEL:-INFO}
+
+LOCUST_OPTS="-f ${LOCUST_FILE} --no-reset-stats"
+
+if [ -z ${HOST_URL+x} ] ; then
+    echo "No value set for (HOST_URL), falling back to host value in the locust class"
+else
+    echo "(HOST_URL) set to ${HOST_URL}"
+    LOCUST_OPTS="--host=${HOST_URL} $LOCUST_OPTS"
 fi
 
-if [ ${1} = "locust" ] ; then 
-    echo "Starting locust......."
-    if [ -z ${LOCUST_HOST+x} ] ; then
-        echo "You need to set the URL of the host to be tested (LOCUST_HOST)."
+case `echo ${LOCUST_MODE} | tr 'a-z' 'A-Z'` in
+"MASTER")
+    LOCUST_OPTS="--master --master-bind-port=${LOCUST_MASTER_BIND_PORT} $LOCUST_OPTS"
+    ;;
+
+"SLAVE")
+    LOCUST_OPTS="--slave --master-host=${LOCUST_MASTER} --master-port=${LOCUST_MASTER_BIND_PORT} $LOCUST_OPTS"
+    if [ -z ${LOCUST_MASTER+x} ] ; then
+        echo "You need to set LOCUST_MASTER."
         exit 1
     fi
-fi
-exec "$@"
\ No newline at end of file
+    ;;
+esac
+
+exec "$@" ${LOCUST_OPTS}
\ No newline at end of file