From eecc604f244a3327dc12f2eb7de2bfdbb70ee4fb Mon Sep 17 00:00:00 2001
From: Alis Akers <alismxwll@gmail.com>
Date: Tue, 10 Dec 2024 11:36:44 -0800
Subject: [PATCH] add a load testing script for process-zip

---
 .../seed-scripts/docker-compose-load.yml      | 32 +++++++
 .../ecr-viewer/seed-scripts/locustfile.py     | 93 +++++++++++++++++++
 2 files changed, 125 insertions(+)
 create mode 100644 containers/ecr-viewer/seed-scripts/docker-compose-load.yml
 create mode 100644 containers/ecr-viewer/seed-scripts/locustfile.py

diff --git a/containers/ecr-viewer/seed-scripts/docker-compose-load.yml b/containers/ecr-viewer/seed-scripts/docker-compose-load.yml
new file mode 100644
index 0000000000..8b5c73b6c4
--- /dev/null
+++ b/containers/ecr-viewer/seed-scripts/docker-compose-load.yml
@@ -0,0 +1,32 @@
+services:
+  # The Locust engine is used for load testing purposes.  
+  locust-brain:
+    image: locustio/locust
+    ports:
+     - "8089:8089"
+    volumes:
+      - ./locustfile.py:/home/locust/locustfile.py
+      - ./requirements.txt:/home/locust/requirements.txt
+      - ./baseECR/LA:/home/locust/baseECR/LA
+    command: -f /home/locust/locustfile.py --master -H $LOAD_TESTING_URL
+  locust-worker-1:
+    image: locustio/locust
+    volumes:
+      - ./locustfile.py:/home/locust/locustfile.py
+      - ./requirements.txt:/home/locust/requirements.txt
+      - ./baseECR/LA:/home/locust/baseECR/LA
+    command: -f /home/locust/locustfile.py --worker --master-host locust-brain
+  locust-worker-2:
+    image: locustio/locust
+    volumes:
+      - ./locustfile.py:/home/locust/locustfile.py
+      - ./requirements.txt:/home/locust/requirements.txt
+      - ./baseECR/LA:/home/locust/baseECR/LA
+    command: -f /home/locust/locustfile.py --worker --master-host locust-brain
+  locust-worker-3:
+    image: locustio/locust
+    volumes:
+      - ./locustfile.py:/home/locust/locustfile.py
+      - ./requirements.txt:/home/locust/requirements.txt
+      - ./baseECR/LA:/home/locust/baseECR/LA
+    command: -f /home/locust/locustfile.py --worker --master-host locust-brain
\ No newline at end of file
diff --git a/containers/ecr-viewer/seed-scripts/locustfile.py b/containers/ecr-viewer/seed-scripts/locustfile.py
new file mode 100644
index 0000000000..b590974253
--- /dev/null
+++ b/containers/ecr-viewer/seed-scripts/locustfile.py
@@ -0,0 +1,93 @@
+from locust import HttpUser, task, between
+import subprocess
+import os
+import shutil
+import random
+import json
+
+class EcrViewer(HttpUser):
+    wait_time = between(1, 5)
+
+    @task
+    def ecr_viewer(self):
+        self.client.get(f"/ecr-viewer")
+
+    @task
+    def orchestration(self):
+        self.client.get(f"/orchestration")
+
+    @task
+    def upload_zip(self):
+        for file in get_zipped_files():
+            with open(file, "rb") as opened_file:
+                data = {
+                    "content_type": "application/zip",
+                    "config_file_name": "sample-orchestration-s3-config.json",
+                    "data_type": "zip",
+                    "message_type": "ecr",
+                }
+                # print(f"Uploading {file}")
+                file_tuple = { "upload_file": ( file, opened_file.read(), "application/zip") }
+                response = self.client.post(f"/orchestration/process-zip", data=data, files=file_tuple)
+                self.tasks.append(check_ecr(self, file, response.json()))
+            
+    # Runs at the start of each test. Useful for authentication and setup actions.
+    def on_start(self):
+        subprocess.run(["pip", "install", "-r", "requirements.txt"])
+        pass
+
+# Check the ecr viewer response for eicr_id and view the ecr
+def check_ecr(self, file, response):
+    if "detail" in response:
+        print(f"{file}", response['detail'])
+    if "message" in response:
+        print(f"{file}", response['message'])
+    if "processed_values" not in response:
+        print("No processed_values found in response")
+        return
+    if "parsed_values" not in response["processed_values"]:
+        print("No parsed_values found in response")
+        return
+    if "eicr_id" in response["processed_values"]["parsed_values"]:
+        print(response["processed_values"]["parsed_values"]["eicr_id"])
+        eicr_id = response["processed_values"]["parsed_values"]["eicr_id"]
+        print(f"/ecr-viewer/view-data?id={eicr_id}")
+        response = self.client.get(f"/ecr-viewer/view-data?id={eicr_id}")
+        print(response)
+    else:
+        print("No eicr_id found in response")
+
+# Get all the zipped files in the baseECR folder
+def get_zipped_files():
+    files = []
+    BASEDIR = os.path.dirname(os.path.abspath(__file__))
+    subfolders = ["LA"]
+    for subfolder in subfolders:
+        subfolder_path = os.path.join(BASEDIR, "baseECR", subfolder)
+
+        # Check if the subfolder exists and is a directory
+        if not os.path.isdir(subfolder_path):
+            print(f"{subfolder_path} is not a valid directory.")
+            continue
+
+        # Now iterate through the folders inside each subfolder
+        for folder in os.listdir(subfolder_path):
+            folder_path = os.path.join(subfolder_path, folder)
+
+            # Check if it's a directory
+            if not os.path.isdir(folder_path):
+                continue
+
+            if os.path.exists(os.path.join(folder_path, "CDA_eICR.xml")):
+                random_number = random.randint(1, 30)
+                zipped_file = shutil.make_archive(f"{random_number}", 'zip', folder_path)
+                print(f"Zipped {folder_path} to {zipped_file}")
+                files.append(zipped_file)
+            # If neither `bundle.json` nor `CDA_eICR.xml` exists, skip processing
+            else:
+                print(
+                    f"Neither `bundle.json` nor `CDA_eICR.xml` found in {folder_path}. Skipping."
+                )
+                continue
+
+    return files