test: Print Ray job logs after successful or failed completion

Author: astefanutti

test/e2e/mnist_rayjob_mcad_raycluster_test.go

@@ -250,13 +250,14 @@ func TestMNISTRayJobMCADRayCluster(t *testing.T) {
 	test.Expect(err).NotTo(HaveOccurred())
 	test.T().Logf("Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name)
 
-	test.T().Logf("Waiting for RayJob %s/%s to complete successfully", rayJob.Namespace, rayJob.Name)
-	test.Eventually(RayJob(test, namespace, rayJob.Name), TestTimeoutLong).
-		Should(WithTransform(RayJobStatus, Equal(rayv1alpha1.JobStatusSucceeded)))
-
-	rayJob, err = test.Client().Ray().RayV1alpha1().RayJobs(namespace.Name).Get(test.Ctx(), rayJob.Name, metav1.GetOptions{})
-	test.Expect(err).NotTo(HaveOccurred())
+	test.T().Logf("Waiting for RayJob %s/%s to complete", rayJob.Namespace, rayJob.Name)
+	test.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutLong).
+		Should(WithTransform(RayJobStatus, Satisfy(rayv1alpha1.IsJobTerminal)))
 
 	test.T().Logf("Printing RayJob %s/%s logs", rayJob.Namespace, rayJob.Name)
-	test.T().Log(GetRayJobLogs(test, rayJob))
+	test.T().Log(GetRayJobLogs(test, rayJob.Namespace, rayJob.Name))
+
+	// Assert the Ray job has completed successfully
+	test.Expect(GetRayJob(test, rayJob.Namespace, rayJob.Name)).
+		To(WithTransform(RayJobStatus, Equal(rayv1alpha1.JobStatusSucceeded)))
 }

test/support/ray.go

@@ -21,28 +21,35 @@ import (
 
 	"github.com/onsi/gomega"
 
-	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	rayv1alpha1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1alpha1"
 )
 
 const RayJobDefaultClusterSelectorKey = "ray.io/cluster"
 
-func RayJob(t Test, namespace *corev1.Namespace, name string) func(g gomega.Gomega) *rayv1alpha1.RayJob {
+func RayJob(t Test, namespace, name string) func(g gomega.Gomega) *rayv1alpha1.RayJob {
 	return func(g gomega.Gomega) *rayv1alpha1.RayJob {
-		job, err := t.Client().Ray().RayV1alpha1().RayJobs(namespace.Name).Get(t.Ctx(), name, metav1.GetOptions{})
+		job, err := t.Client().Ray().RayV1alpha1().RayJobs(namespace).Get(t.Ctx(), name, metav1.GetOptions{})
 		g.Expect(err).NotTo(gomega.HaveOccurred())
 		return job
 	}
 }
 
+func GetRayJob(t Test, namespace, name string) *rayv1alpha1.RayJob {
+	t.T().Helper()
+	return RayJob(t, namespace, name)(t)
+}
+
 func RayJobStatus(job *rayv1alpha1.RayJob) rayv1alpha1.JobStatus {
 	return job.Status.JobStatus
 }
 
-func GetRayJobLogs(t Test, job *rayv1alpha1.RayJob) string {
+func GetRayJobLogs(t Test, namespace, name string) string {
 	t.T().Helper()
+
+	job := GetRayJob(t, namespace, name)
+
 	response := t.Client().Core().CoreV1().RESTClient().
 		Get().
 		AbsPath("/api/v1/namespaces", job.Namespace, "services", "http:"+job.Status.RayClusterName+"-head-svc:dashboard", "proxy", "api", "jobs", job.Status.JobId, "logs").