Add common functions for PyTorchJob

Author: sutaakar

support/fakeclient.go

@@ -3,6 +3,7 @@ package support
 import (
 	"testing"
 
+	fakekfto "github.com/kubeflow/training-operator/pkg/client/clientset/versioned/fake"
 	fakeray "github.com/ray-project/kuberay/ray-operator/pkg/client/clientset/versioned/fake"
 
 	"k8s.io/apimachinery/pkg/runtime"
@@ -22,17 +23,19 @@ func NewTest(t *testing.T) *T {
 	fakerouteClient := fakeroute.NewSimpleClientset()
 	fakerayClient := fakeray.NewSimpleClientset()
 	fakekueueClient := fakekueue.NewSimpleClientset()
+	fakeKftoClient := fakekfto.NewSimpleClientset()
 	fakedynamicClient := fakeDynamic.NewSimpleDynamicClient(runtime.NewScheme())
 
 	test := With(t).(*T)
 	test.client = &testClient{
-		core:    fakeCoreClient,
-		machine: fakemachineClient,
-		image:   fakeimageClient,
-		route:   fakerouteClient,
-		ray:     fakerayClient,
-		kueue:   fakekueueClient,
-		dynamic: fakedynamicClient,
+		core:     fakeCoreClient,
+		machine:  fakemachineClient,
+		image:    fakeimageClient,
+		route:    fakerouteClient,
+		ray:      fakerayClient,
+		kueue:    fakekueueClient,
+		kubeflow: fakeKftoClient,
+		dynamic:  fakedynamicClient,
 	}
 	return test
 }

support/pytorchjob.go

@@ -0,0 +1,66 @@
+/*
+Copyright 2024.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package support
+
+import (
+	kftov1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
+	"github.com/onsi/gomega"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func PyTorchJob(t Test, namespace, name string) func(g gomega.Gomega) *kftov1.PyTorchJob {
+	return func(g gomega.Gomega) *kftov1.PyTorchJob {
+		job, err := t.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Get(t.Ctx(), name, metav1.GetOptions{})
+		g.Expect(err).NotTo(gomega.HaveOccurred())
+		return job
+	}
+}
+
+func PyTorchJobs(t Test, namespace string) func(g gomega.Gomega) []kftov1.PyTorchJob {
+	return func(g gomega.Gomega) []kftov1.PyTorchJob {
+		jobs, err := t.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).List(t.Ctx(), metav1.ListOptions{})
+		g.Expect(err).NotTo(gomega.HaveOccurred())
+		return jobs.Items
+	}
+}
+
+func PyTorchJobConditionRunning(job *kftov1.PyTorchJob) corev1.ConditionStatus {
+	return PyTorchJobCondition(job, kftov1.JobRunning)
+}
+
+func PyTorchJobConditionSucceeded(job *kftov1.PyTorchJob) corev1.ConditionStatus {
+	return PyTorchJobCondition(job, kftov1.JobSucceeded)
+}
+
+func PyTorchJobConditionSuspended(job *kftov1.PyTorchJob) corev1.ConditionStatus {
+	return PyTorchJobCondition(job, kftov1.JobSuspended)
+}
+
+func PyTorchJobConditionFailed(job *kftov1.PyTorchJob) corev1.ConditionStatus {
+	return PyTorchJobCondition(job, kftov1.JobFailed)
+}
+
+func PyTorchJobCondition(job *kftov1.PyTorchJob, conditionType kftov1.JobConditionType) corev1.ConditionStatus {
+	for _, condition := range job.Status.Conditions {
+		if condition.Type == conditionType {
+			return condition.Status
+		}
+	}
+	return corev1.ConditionUnknown
+}

support/pytorchjob_test.go

@@ -0,0 +1,73 @@
+/*
+Copyright 2024.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package support
+
+import (
+	"testing"
+
+	kftov1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
+	"github.com/onsi/gomega"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestPyTorchJob(t *testing.T) {
+
+	test := NewTest(t)
+
+	job := &kftov1.PyTorchJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "my-job-1",
+			Namespace: "my-namespace",
+		},
+	}
+
+	test.client.Kubeflow().KubeflowV1().PyTorchJobs("my-namespace").Create(test.ctx, job, metav1.CreateOptions{})
+
+	pyTorchJob := PyTorchJob(test, "my-namespace", "my-job-1")(test)
+	test.Expect(pyTorchJob.Name).To(gomega.Equal("my-job-1"))
+	test.Expect(pyTorchJob.Namespace).To(gomega.Equal("my-namespace"))
+}
+
+func TestPyTorchJobs(t *testing.T) {
+
+	test := NewTest(t)
+
+	jobOne := &kftov1.PyTorchJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "my-job-1",
+			Namespace: "my-namespace",
+		},
+	}
+	jobTwo := &kftov1.PyTorchJob{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "my-job-2",
+			Namespace: "my-namespace",
+		},
+	}
+
+	test.client.Kubeflow().KubeflowV1().PyTorchJobs("my-namespace").Create(test.ctx, jobOne, metav1.CreateOptions{})
+	test.client.Kubeflow().KubeflowV1().PyTorchJobs("my-namespace").Create(test.ctx, jobTwo, metav1.CreateOptions{})
+
+	pyTorchJobs := PyTorchJobs(test, "my-namespace")(test)
+
+	test.Expect(len(pyTorchJobs)).To(gomega.Equal(2))
+	test.Expect(pyTorchJobs[0].Name).To(gomega.Equal("my-job-1"))
+	test.Expect(pyTorchJobs[0].Namespace).To(gomega.Equal("my-namespace"))
+	test.Expect(pyTorchJobs[1].Name).To(gomega.Equal("my-job-2"))
+	test.Expect(pyTorchJobs[1].Namespace).To(gomega.Equal("my-namespace"))
+}