From a3200c95e787086e0778df9a37669160c33f9298 Mon Sep 17 00:00:00 2001
From: Jiri Petrlik <jiripetrlik@gmail.com>
Date: Thu, 9 May 2024 23:09:25 +0200
Subject: [PATCH] RHOAIENG-5344 - E2E test for Ray local interactive

---
 tests/e2e/local_interactive_sdk_kind_test.py  | 71 +++++++++++++++
 tests/e2e/local_interactive_sdk_oauth_test.py | 86 +++++++++++++++++++
 2 files changed, 157 insertions(+)
 create mode 100644 tests/e2e/local_interactive_sdk_kind_test.py
 create mode 100644 tests/e2e/local_interactive_sdk_oauth_test.py

diff --git a/tests/e2e/local_interactive_sdk_kind_test.py b/tests/e2e/local_interactive_sdk_kind_test.py
new file mode 100644
index 000000000..efc6c8f3e
--- /dev/null
+++ b/tests/e2e/local_interactive_sdk_kind_test.py
@@ -0,0 +1,71 @@
+from codeflare_sdk import Cluster, ClusterConfiguration
+
+import pytest
+import ray
+import math
+
+from support import *
+
+
+@pytest.mark.kind
+class TestRayLocalInteractiveOauth:
+    def setup_method(self):
+        initialize_kubernetes_client(self)
+
+    def teardown_method(self):
+        delete_namespace(self)
+
+    def test_local_interactives(self):
+        self.setup_method()
+        create_namespace(self)
+        create_kueue_resources(self)
+        self.run_local_interactives()
+
+    def run_local_interactives(self):
+        ray_image = get_ray_image()
+
+        cluster_name = "test-ray-cluster"
+
+        cluster = Cluster(
+            ClusterConfiguration(
+                namespace=self.namespace,
+                name=cluster_name,
+                num_workers=1,
+                min_cpus=1,
+                max_cpus=1,
+                min_memory=4,
+                max_memory=4,
+                num_gpus=0,
+                image=ray_image,
+                verify_tls=False,
+            )
+        )
+        cluster.up()
+        cluster.wait_ready()
+
+        ray.shutdown()
+        ray.init(address=cluster.local_client_url(), logging_level="DEBUG")
+
+        @ray.remote
+        def heavy_calculation_part(num_iterations):
+            result = 0.0
+            for i in range(num_iterations):
+                for j in range(num_iterations):
+                    for k in range(num_iterations):
+                        result += math.sin(i) * math.cos(j) * math.tan(k)
+            return result
+
+        @ray.remote
+        def heavy_calculation(num_iterations):
+            results = ray.get(
+                [heavy_calculation_part.remote(num_iterations // 30) for _ in range(30)]
+            )
+            return sum(results)
+
+        ref = heavy_calculation.remote(3000)
+        result = ray.get(ref)
+        assert result == 1789.4644387076714
+        ray.cancel(ref)
+        ray.shutdown()
+
+        cluster.down()
diff --git a/tests/e2e/local_interactive_sdk_oauth_test.py b/tests/e2e/local_interactive_sdk_oauth_test.py
new file mode 100644
index 000000000..f121038d7
--- /dev/null
+++ b/tests/e2e/local_interactive_sdk_oauth_test.py
@@ -0,0 +1,86 @@
+from codeflare_sdk import (
+    Cluster,
+    ClusterConfiguration,
+    TokenAuthentication,
+    generate_cert,
+)
+
+import math
+import pytest
+import ray
+
+from support import *
+
+
+@pytest.mark.openshift
+class TestRayLocalInteractiveOauth:
+    def setup_method(self):
+        initialize_kubernetes_client(self)
+
+    def teardown_method(self):
+        delete_namespace(self)
+
+    def test_local_interactives(self):
+        self.setup_method()
+        create_namespace(self)
+        create_kueue_resources(self)
+        self.run_local_interactives()
+
+    def run_local_interactives(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_name = "test-ray-cluster"
+
+        cluster = Cluster(
+            ClusterConfiguration(
+                namespace=self.namespace,
+                name=cluster_name,
+                num_workers=1,
+                min_cpus=1,
+                max_cpus=1,
+                min_memory=4,
+                max_memory=4,
+                num_gpus=0,
+                image=ray_image,
+                verify_tls=False,
+            )
+        )
+        cluster.up()
+        cluster.wait_ready()
+
+        generate_cert.generate_tls_cert(cluster_name, self.namespace)
+        generate_cert.export_env(cluster_name, self.namespace)
+
+        ray.shutdown()
+        ray.init(address=cluster.local_client_url(), logging_level="DEBUG")
+
+        @ray.remote
+        def heavy_calculation_part(num_iterations):
+            result = 0.0
+            for i in range(num_iterations):
+                for j in range(num_iterations):
+                    for k in range(num_iterations):
+                        result += math.sin(i) * math.cos(j) * math.tan(k)
+            return result
+
+        @ray.remote
+        def heavy_calculation(num_iterations):
+            results = ray.get(
+                [heavy_calculation_part.remote(num_iterations // 30) for _ in range(30)]
+            )
+            return sum(results)
+
+        ref = heavy_calculation.remote(3000)
+        result = ray.get(ref)
+        assert result == 1789.4644387076714
+        ray.cancel(ref)
+        ray.shutdown()
+
+        cluster.down()