From 18548abc6905ae81307a6da0ba7141836a7418a9 Mon Sep 17 00:00:00 2001
From: Srihari <svenkata@redhat.com>
Date: Thu, 21 Mar 2024 22:23:31 +0530
Subject: [PATCH] Add TokenAuthentication step for RayCluster oauth test

---
 tests/e2e/mnist_raycluster_sdk_oauth_test.py |  9 ++++++++-
 tests/e2e/support.py                         | 12 ++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tests/e2e/mnist_raycluster_sdk_oauth_test.py b/tests/e2e/mnist_raycluster_sdk_oauth_test.py
index 07991a30..3e24d465 100644
--- a/tests/e2e/mnist_raycluster_sdk_oauth_test.py
+++ b/tests/e2e/mnist_raycluster_sdk_oauth_test.py
@@ -4,7 +4,7 @@
 
 from torchx.specs.api import AppState, is_terminal
 
-from codeflare_sdk.cluster.cluster import Cluster, ClusterConfiguration
+from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
 from codeflare_sdk.job.jobs import DDPJobDefinition
 
 import pytest
@@ -30,6 +30,13 @@ def test_mnist_ray_cluster_sdk_auth(self):
     def run_mnist_raycluster_sdk_oauth(self):
         ray_image = get_ray_image()
 
+        auth = TokenAuthentication(
+            token=run_oc_command(["whoami", "--show-token=true"]),
+            server=run_oc_command(["whoami", "--show-server=true"]),
+            skip_tls=True,
+        )
+        auth.login()
+
         cluster = Cluster(
             ClusterConfiguration(
                 name="mnist",
diff --git a/tests/e2e/support.py b/tests/e2e/support.py
index 1ab45c19..abf4ac66 100644
--- a/tests/e2e/support.py
+++ b/tests/e2e/support.py
@@ -1,6 +1,7 @@
 import os
 import random
 import string
+import subprocess
 from kubernetes import client, config
 import kubernetes.client
 
@@ -33,3 +34,14 @@ def initialize_kubernetes_client(self):
     # Initialize Kubernetes client
     self.api_instance = client.CoreV1Api()
     self.custom_api = kubernetes.client.CustomObjectsApi(self.api_instance.api_client)
+
+
+def run_oc_command(args):
+    try:
+        result = subprocess.run(
+            ["oc"] + args, capture_output=True, text=True, check=True
+        )
+        return result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        print(f"Error executing 'oc {' '.join(args)}': {e}")
+        return None