Workspaces Launch Test

files/scripts/workspaces_launch_test.py

@@ -1,17 +1,15 @@
 import time
 import argparse
-import json
-import jwt
 import logging
 
 import requests
 
-COMMON_URL = "https://binamb.planx-pla.net"
-TOKEN_REFRESH_THRESHOLD = 60
+workspace_internal_url = "http://workspace-token-service"
 logging.basicConfig(level=logging.DEBUG)
 
 def main():
-    tester = WorkspaceLaunchTest(commons_url=COMMON_URL, api_file="/Users/binambajracharya/Downloads/credentials.json")
+    args = parse_args()
+    tester = WorkspaceLaunchTest(commons_url=args.COMMONS_URL, access_token=args.access_token)
     tester.start_workspace_launch_test()
 
 def parse_args():
@@ -24,97 +22,34 @@ def parse_args():
         help="Specify the Commons URL to test"
     )
     parser.add_argument(
-        "--api-file",
-        dest="api_file",
-        help="API Key file",
+        "--access-token",
+        dest="access_token",
+        help="User's access token. It should have the 'credentials' scope since the /launch api requires an access token that can use to get an api key.",
     )
 
     return parser.parse_args()
 
 
 class WorkspaceLaunchTest:
-    def __init__(self, commons_url, api_file):
+    def __init__(self, commons_url, access_token):
         self.commons_url = commons_url
-        self.api_file = api_file
-        self.access_token = None
+        self.workspace_internal_url = workspace_internal_url
         self.token_expiration = 0
         self.headers = {}
         self.start_time = 0
-        if self.api_file:
-            self.api_file = self.get_api_key_from_file(self.api_file)
-            self.access_token, self.token_expiration = self.get_access_token_with_key(self.api_file)
-            self.update_headers()
-
-
-    def get_api_key_from_file(self, filepath):
-        try:
-            with open(filepath, "r") as f:
-                data = json.load(f)
-
-            api_key = data.get("api_key")
-
-            if api_key:
-                return api_key
-            else:
-                logging.error(f"Could not get API Key from json file: {filepath}")
-        except FileNotFoundError:
-            logging.error(f"Could not find file: {filepath}")
-            return None
-        except Exception as e:
-            logging.error(f"Could not get API Key from json file with error : {e}")
-
-
-    def get_access_token_with_key(self, api_key):
-        """
-        Try to fetch an access token given the api key
-        """
-        # attempt to get a token from Fence
-        body = {"scope": [
-                    "openid",
-                    "user",
-                    "data",
-                    "credentials",
-
-                    "ga4gh_passport_v1",
-                    "google_link",
-                    "google_credentials",
-                    "admin",
-                    "fence",
-                    "google_service_account"
-                    ],
-                "api_key": api_key
-                }
-        auth_url = "{}/user/credentials/api/access_token".format(COMMON_URL)
-        try:
-            resp = requests.post(auth_url, json=body)
-            resp.raise_for_status()
-            access_token = resp.json()["access_token"]
-            decoded_token = jwt.decode(
-                access_token, algorithms=["RS256"], options={"verify_signature": False}
-                )
-            print(access_token)
-            token_expiration = decoded_token["exp"]
-            logging.info("Successfully retrieved access token.")
-            return access_token, token_expiration
-        except requests.exceptions.RequestException as e:
-            logging.error(f"Error fetching access token: {e}")
-            return None, 0
-
+        self.access_token = access_token
+        self.launch_status = "Workspace did not launch. Something went wrong before launch."
+        self.reason_for_failure = None
+        self.update_headers()
+
+        self.json_result = {
+            "start_time": int,
+            "end_time": int,
+            "duration": int,
+            "result": str,
+            "reason_for_failure": str
+        }
 
-    def refresh_access_token_if_needed(self):
-        """
-        Refresh the access token if its about to expire. 
-        """
-        if not self.api_file:
-            return
-
-        if time.time() > self.token_expiration - TOKEN_REFRESH_THRESHOLD:
-            logging.info("Access token is about to expire. Refreshing token...")
-            self.access_token, self.token_expiration = self.get_access_token_with_key(self.api_file)
-            self.update_headers()
-            if not self.access_token:
-                logging.error("Failed to refresh access token")
-
     def update_headers(self):
         """Updates the headers with the current access token."""
         if self.access_token:
@@ -128,14 +63,15 @@ def start_workspace_launch_test(self):
             "Authorization": f"Bearer {self.access_token}"
         }
         # Get available workspace options
-        options_url = COMMON_URL + "/lw-workspace/options"
+        options_url = self.commons_url + "/lw-workspace/options"
         try:
-            self.refresh_access_token_if_needed()
             options_response = requests.get(options_url, headers=self.headers)
             options_response.raise_for_status()
 
         except requests.exceptions.RequestException as e:
-            logging.error(f"Couldn't get workspace options with error: {e}")
+            error_msg = f"Couldn't get workspace options with error: {e}"
+            logging.error(error_msg)
+            self.reason_for_failure = error_msg
 
         options = options_response.json()
         logging.info("Successfully found workspace options")
@@ -144,14 +80,14 @@ def start_workspace_launch_test(self):
         workspace_id = options[0].get("id")
 
         # Launch workspace
-        launch_url = COMMON_URL + "/lw-workspace/launch" + "?id=" + workspace_id
+        launch_url = self.commons_url + "/lw-workspace/launch" + "?id=" + workspace_id
         try:
-            self.refresh_access_token_if_needed()
             launch_response = requests.post(launch_url, headers=self.headers)
-            launch_response.raise_for_status()
             self.start_time = time.time()
         except requests.exceptions.RequestException as e:
-            logging.error(f"Couldn't launch workspace. Error code with error: {e}")
+            error_msg = f"Couldn't launch workspace. Error code with error: {e}"
+            logging.error(error_msg)
+            self.reason_for_failure = error_msg
             return
 
         logging.info("Successfully started launching workspace. Starting timer and monitoring workspace status...")
@@ -162,38 +98,49 @@ def start_workspace_launch_test(self):
         logging.info(f"Workspace took {end_time-self.start_time} seconds to initialize")
 
         # Terminate active running workspace
-        terminate_url = COMMON_URL + "/lw-workspace/terminate"
+        terminate_url = self.commons_url + "/lw-workspace/terminate"
         try:
-            self.refresh_access_token_if_needed()
             logging.info("Attempting to terminate workspace...")
             terminate_response = requests.post(terminate_url, headers= self.headers)
             terminate_response.raise_for_status()
             logging.info("Workspace terminated...")
         except requests.exceptions.RequestException as e:
-            logging.error(f"Couldn't terminate workspace with error : {e}")
+            error_msg = f"Couldn't terminate workspace with error : {e}"
+            logging.error(error_msg)
+            self.reason_for_failure = error_msg
 
         logging.info("Workspace terminated...")
+
+        logging.info("Result: ")
+        self.json_result["start_time"] = self.start_time
+        self.json_result["end_time"] = end_time
+        self.json_result["duration"] = end_time - self.start_time
+        self.json_result["result"] = self.launch_status
+        self.json_result["reason_for_failure"] = self.reason_for_failure
 
+        logging.info(self.json_result)
 
     def monitor_workspace_status(self, interval=10):
-        status_url = COMMON_URL + "/lw-workspace/status"
+        status_url = self.commons_url + "/lw-workspace/status"
 
         while True:
             try:
                 status_response = requests.get(status_url, headers=self.headers)
                 status_response.raise_for_status()
             except requests.exceptions.RequestException as e:
-                logging.error(f"Error checking workspace status: {e}")
+                error_msg = f"Error checking workspace status: {e}"
+                logging.error(error_msg)
+                self.reason_for_failure = error_msg
 
             logging.info(f"Status reposnse: {status_response.json()}")
 
             if status_response.json()["status"] == "Running":
+                self.launch_status = "Running"
                 break
 
             time.sleep(interval)
             logging.info(f"Elapsed time: {time.time()-self.start_time}")
-
-
+            self.launch_status = status_response.json()["status"]
 
 
 if __name__ == "__main__":

kube/services/jobs/workspace-launch-test-cronjob.yaml

@@ -0,0 +1,34 @@
+apiVersion: batch/v1
+kind: CronJob
+metadata:
+  name: workspace-launch-test
+spec:
+  schedule: 60 * * * *
+  successfulJobsHistoryLimit: 2
+  failedJobsHistoryLimit: 2
+  concurrencyPolicy: Forbid
+  jobTemplate:
+    spec:
+      backoffLimit: 0 
+      template:
+        metadata:
+          labels:
+            app:gen3job
+        spec:
+          affinity:
+            nodeAffinity:
+              preferredDuringSchedulingIgnoredDuringExecution:
+              - weight: 100
+                preference:
+                  matchExpressions:
+                  - key: karpenter.sh/capacity-type
+                    operator: In
+                    values:
+                    - on-demand
+              - weight: 99
+                preference:
+                  matchExpressions:
+                  - key: eks.amazonaws.com/capacityType
+                    operator: In
+                    values:
+                    - ONDEMAND