From 1d5a8242d0f5784ff72102f5d7d549b3a6eb853e Mon Sep 17 00:00:00 2001
From: Stephen Astels <sastels@gmail.com>
Date: Tue, 24 Oct 2023 13:58:12 -0400
Subject: [PATCH] soak test with sms

---
 scripts/soak_test/.env.example                |  1 +
 scripts/soak_test/soak_test_send_email.py     | 30 ---------------
 .../soak_test/soak_test_send_notification.py  | 38 +++++++++++++++++++
 3 files changed, 39 insertions(+), 30 deletions(-)
 delete mode 100644 scripts/soak_test/soak_test_send_email.py
 create mode 100644 scripts/soak_test/soak_test_send_notification.py

diff --git a/scripts/soak_test/.env.example b/scripts/soak_test/.env.example
index 0da06802cc..6e797ff095 100644
--- a/scripts/soak_test/.env.example
+++ b/scripts/soak_test/.env.example
@@ -1,2 +1,3 @@
 API_KEY=
 EMAIL_TEMPLATE_ID=
+SMS_TEMPLATE_ID=
diff --git a/scripts/soak_test/soak_test_send_email.py b/scripts/soak_test/soak_test_send_email.py
deleted file mode 100644
index f9cefd88d1..0000000000
--- a/scripts/soak_test/soak_test_send_email.py
+++ /dev/null
@@ -1,30 +0,0 @@
-import os
-
-from dotenv import load_dotenv
-from locust import HttpUser, constant_pacing, events, task
-from soak_utils import url_with_prefix
-
-load_dotenv()
-
-
-@events.init_command_line_parser.add_listener
-def _(parser):
-    parser.add_argument("--ref", type=str, default="test", help="reference")
-
-
-class NotifyApiUser(HttpUser):
-    wait_time = constant_pacing(1)  # each user makes one post per second
-
-    def __init__(self, *args, **kwargs):
-        self.host = url_with_prefix(self.host, "api")
-
-        super(NotifyApiUser, self).__init__(*args, **kwargs)
-        self.headers = {"Authorization": f"apikey-v1 {os.getenv('API_KEY')}"}
-        self.email_address = "success@simulator.amazonses.com"
-        self.email_template = os.getenv("EMAIL_TEMPLATE_ID")
-
-    @task(1)
-    def send_email(self):
-        reference_id = self.environment.parsed_options.ref
-        json = {"email_address": self.email_address, "template_id": self.email_template, "reference": reference_id}
-        self.client.post("/v2/notifications/email", json=json, headers=self.headers)
diff --git a/scripts/soak_test/soak_test_send_notification.py b/scripts/soak_test/soak_test_send_notification.py
new file mode 100644
index 0000000000..be8f3b0ece
--- /dev/null
+++ b/scripts/soak_test/soak_test_send_notification.py
@@ -0,0 +1,38 @@
+import os
+
+from dotenv import load_dotenv
+from locust import HttpUser, constant_pacing, events, task
+from soak_utils import url_with_prefix
+
+load_dotenv()
+
+
+@events.init_command_line_parser.add_listener
+def _(parser):
+    parser.add_argument("--ref", type=str, default="test", help="reference")
+    parser.add_argument("--type", type=str, default="email", help="email or sms")
+
+
+class NotifyApiUser(HttpUser):
+    wait_time = constant_pacing(1)  # each user makes one post per second
+
+    def __init__(self, *args, **kwargs):
+        self.host = url_with_prefix(self.host, "api")
+
+        super(NotifyApiUser, self).__init__(*args, **kwargs)
+        self.headers = {"Authorization": f"apikey-v1 {os.getenv('API_KEY')}"}
+        self.email_template = os.getenv("EMAIL_TEMPLATE_ID")
+        self.sms_template = os.getenv("SMS_TEMPLATE_ID")
+        self.email_address = "success@simulator.amazonses.com"
+        self.phone_number = "+16135550123" # INTERNAL_TEST_NUMBER
+        self.reference_id = self.environment.parsed_options.ref
+        self.type = self.environment.parsed_options.type
+
+    @task(1)
+    def send_notification(self):
+        if self.type == "email":
+            json = {"email_address": self.email_address, "template_id": self.email_template, "reference": self.reference_id}
+            self.client.post("/v2/notifications/email", json=json, headers=self.headers)
+        else:
+            json = {"phone_number": self.phone_number, "template_id": self.sms_template, "reference": self.reference_id}
+            self.client.post("/v2/notifications/sms", json=json, headers=self.headers)
\ No newline at end of file