Add kueue integration tests for Trainer v2

Author: ChughShilpa

go.mod

@@ -2,20 +2,28 @@ module github.com/opendatahub-io/distributed-workloads
 
 go 1.24.2
 
+// Remove replace when Training operator v1 tests are removed
+replace (
+	github.com/kubeflow/training-operator => github.com/kubeflow/training-operator v1.7.0
+	sigs.k8s.io/kueue => sigs.k8s.io/kueue v0.6.2
+)
+
 require (
 	github.com/kubeflow/trainer/v2 v2.0.0
 	github.com/kubeflow/training-operator v1.7.0
 	github.com/matoous/go-nanoid/v2 v2.1.0
-	github.com/onsi/gomega v1.37.0
-	github.com/openshift/api v0.0.0-20251015095338-264e80a2b6e7
+	github.com/onsi/gomega v1.38.2
+	github.com/openshift/api v0.0.0-20251124165233-999c45c0835a
 	github.com/openshift/client-go v0.0.0-20251015124057-db0dee36e235
-	github.com/prometheus/client_golang v1.22.0
-	github.com/prometheus/common v0.62.0
+	github.com/openshift/kueue-operator v0.0.0-20251202204851-958c48004dad
+	github.com/prometheus/client_golang v1.23.0
+	github.com/prometheus/common v0.65.0
 	github.com/ray-project/kuberay/ray-operator v1.3.0
 	k8s.io/api v0.34.1
 	k8s.io/apimachinery v0.34.1
 	k8s.io/client-go v0.34.1
-	sigs.k8s.io/kueue v0.6.2
+	sigs.k8s.io/jobset v0.9.1
+	sigs.k8s.io/kueue v0.10.2
 )
 
 require (
@@ -25,9 +33,9 @@ require (
 	github.com/emicklei/go-restful/v3 v3.12.2 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
 	github.com/go-logr/logr v1.4.3 // indirect
-	github.com/go-openapi/jsonpointer v0.21.0 // indirect
+	github.com/go-openapi/jsonpointer v0.21.1 // indirect
 	github.com/go-openapi/jsonreference v0.21.0 // indirect
-	github.com/go-openapi/swag v0.23.0 // indirect
+	github.com/go-openapi/swag v0.23.1 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/google/gnostic-models v0.7.0 // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
@@ -41,23 +49,20 @@ require (
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/prometheus/client_model v0.6.1 // indirect
-	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/prometheus/client_model v0.6.2 // indirect
+	github.com/prometheus/procfs v0.16.1 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/spf13/pflag v1.0.6 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
-	go.etcd.io/etcd/client/pkg/v3 v3.5.21 // indirect
-	go.etcd.io/etcd/client/v3 v3.5.21 // indirect
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.58.0 // indirect
 	go.yaml.in/yaml/v2 v2.4.2 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect
-	golang.org/x/net v0.38.0 // indirect
-	golang.org/x/oauth2 v0.27.0 // indirect
-	golang.org/x/sys v0.32.0 // indirect
-	golang.org/x/term v0.30.0 // indirect
-	golang.org/x/text v0.23.0 // indirect
-	golang.org/x/time v0.10.0 // indirect
-	google.golang.org/protobuf v1.36.5 // indirect
+	golang.org/x/net v0.47.0 // indirect
+	golang.org/x/oauth2 v0.30.0 // indirect
+	golang.org/x/sys v0.38.0 // indirect
+	golang.org/x/term v0.37.0 // indirect
+	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/time v0.11.0 // indirect
+	google.golang.org/protobuf v1.36.7 // indirect
 	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
@@ -66,10 +71,8 @@ require (
 	k8s.io/kube-openapi v0.0.0-20250710124328-f3f2b991d03b // indirect
 	k8s.io/utils v0.0.0-20251002143259-bc988d571ff4 // indirect
 	sigs.k8s.io/controller-runtime v0.21.0 // indirect
-	sigs.k8s.io/jobset v0.8.2 // indirect
 	sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8 // indirect
 	sigs.k8s.io/randfill v1.0.0 // indirect
-	sigs.k8s.io/structured-merge-diff/v4 v4.7.0 // indirect
 	sigs.k8s.io/structured-merge-diff/v6 v6.3.0 // indirect
 	sigs.k8s.io/yaml v1.6.0 // indirect
 )

go.sum

@@ -33,6 +33,7 @@ import (
 	imagev1 "github.com/openshift/client-go/image/clientset/versioned"
 	machinev1 "github.com/openshift/client-go/machine/clientset/versioned"
 	routev1 "github.com/openshift/client-go/route/clientset/versioned"
+	kueueoperatorclient "github.com/openshift/kueue-operator/pkg/generated/clientset/versioned"
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
 	// to ensure that exec-entrypoint and run can make use of them.
 )
@@ -43,6 +44,7 @@ type Client interface {
 	Kubeflow() kubeflowclient.Interface
 	JobSet() jobsetclient.Interface
 	Kueue() kueueclient.Interface
+	KueueOperator() kueueoperatorclient.Interface
 	Machine() machinev1.Interface
 	Route() routev1.Interface
 	Image() imagev1.Interface
@@ -52,17 +54,18 @@ type Client interface {
 }
 
 type testClient struct {
-	core     kubernetes.Interface
-	trainer  trainerclient.Interface
-	kubeflow kubeflowclient.Interface
-	jobset   jobsetclient.Interface
-	kueue    kueueclient.Interface
-	machine  machinev1.Interface
-	route    routev1.Interface
-	image    imagev1.Interface
-	ray      rayclient.Interface
-	dynamic  dynamic.Interface
-	storage  storageclient.StorageV1Interface
+	core          kubernetes.Interface
+	trainer       trainerclient.Interface
+	kubeflow      kubeflowclient.Interface
+	jobset        jobsetclient.Interface
+	kueue         kueueclient.Interface
+	kueueOperator kueueoperatorclient.Interface
+	machine       machinev1.Interface
+	route         routev1.Interface
+	image         imagev1.Interface
+	ray           rayclient.Interface
+	dynamic       dynamic.Interface
+	storage       storageclient.StorageV1Interface
 }
 
 var _ Client = (*testClient)(nil)
@@ -87,6 +90,10 @@ func (t *testClient) Kueue() kueueclient.Interface {
 	return t.kueue
 }
 
+func (t *testClient) KueueOperator() kueueoperatorclient.Interface {
+	return t.kueueOperator
+}
+
 func (t *testClient) Machine() machinev1.Interface {
 	return t.machine
 }
@@ -148,6 +155,11 @@ func newTestClient(cfg *rest.Config) (Client, *rest.Config, error) {
 		return nil, nil, err
 	}
 
+	kueueOperatorClient, err := kueueoperatorclient.NewForConfig(cfg)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	machineClient, err := machinev1.NewForConfig(cfg)
 	if err != nil {
 		return nil, nil, err
@@ -179,16 +191,17 @@ func newTestClient(cfg *rest.Config) (Client, *rest.Config, error) {
 	}
 
 	return &testClient{
-		core:     kubeClient,
-		trainer:  trainerClient,
-		kubeflow: kubeflowClient,
-		jobset:   jobsetClient,
-		kueue:    kueueClient,
-		machine:  machineClient,
-		route:    routeClient,
-		image:    imageClient,
-		ray:      rayClient,
-		dynamic:  dynamicClient,
-		storage:  storageClient,
+		core:          kubeClient,
+		trainer:       trainerClient,
+		kubeflow:      kubeflowClient,
+		jobset:        jobsetClient,
+		kueue:         kueueClient,
+		kueueOperator: kueueOperatorClient,
+		machine:       machineClient,
+		route:         routeClient,
+		image:         imageClient,
+		ray:           rayClient,
+		dynamic:       dynamicClient,
+		storage:       storageClient,
 	}, cfg, nil
 }

tests/common/support/client.go

@@ -0,0 +1,125 @@
+/*
+Copyright 2025.
+
+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 (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/onsi/gomega"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/client-go/dynamic"
+)
+
+var DscGVR = schema.GroupVersionResource{
+	Group:    "datasciencecluster.opendatahub.io",
+	Version:  "v2",
+	Resource: "datascienceclusters",
+}
+
+func getDSC(dynamicClient dynamic.Interface, ctx context.Context, name string) (*unstructured.Unstructured, error) {
+	return dynamicClient.Resource(DscGVR).Get(ctx, name, metav1.GetOptions{})
+}
+
+func updateDSC(dynamicClient dynamic.Interface, ctx context.Context, dsc *unstructured.Unstructured) (*unstructured.Unstructured, error) {
+	return dynamicClient.Resource(DscGVR).Update(ctx, dsc, metav1.UpdateOptions{})
+}
+
+func setComponentState(dynamicClient dynamic.Interface, ctx context.Context, dscName, component, state string) error {
+	dsc, err := getDSC(dynamicClient, ctx, dscName)
+	if err != nil {
+		return err
+	}
+
+	err = unstructured.SetNestedField(dsc.Object, state, "spec", "components", component, "managementState")
+	if err != nil {
+		return err
+	}
+
+	_, err = updateDSC(dynamicClient, ctx, dsc)
+	return err
+}
+
+func WaitForComponentState(dynamicClient dynamic.Interface, ctx context.Context, dscName, component, expectedState string, timeout time.Duration) error {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		dsc, err := getDSC(dynamicClient, ctx, dscName)
+		if err != nil {
+			return err
+		}
+		currentState := ComponentStatusManagementState(dsc, component)
+		if currentState == expectedState {
+			return nil
+		}
+		time.Sleep(2 * time.Second)
+	}
+	return fmt.Errorf("timeout waiting for component %s to reach state %s", component, expectedState)
+}
+
+func SetComponentStateAndWait(dynamicClient dynamic.Interface, ctx context.Context, dscName, component, state string, timeout time.Duration) error {
+	if err := setComponentState(dynamicClient, ctx, dscName, component, state); err != nil {
+		return err
+	}
+	return WaitForComponentState(dynamicClient, ctx, dscName, component, state, timeout)
+}
+
+func GetDSC(test Test, name string) (*unstructured.Unstructured, error) {
+	return getDSC(test.Client().Dynamic(), test.Ctx(), name)
+}
+
+func SetComponentToUnmanaged(test Test, dscName string, component string) error {
+	return setComponentState(test.Client().Dynamic(), test.Ctx(), dscName, component, "Unmanaged")
+}
+
+func DSCResource(test Test, name string) func(g gomega.Gomega) *unstructured.Unstructured {
+	return func(g gomega.Gomega) *unstructured.Unstructured {
+		dsc, err := GetDSC(test, name)
+		g.Expect(err).NotTo(gomega.HaveOccurred())
+		return dsc
+	}
+}
+
+func ComponentStatusManagementState(dsc *unstructured.Unstructured, component string) string {
+	state, found, err := unstructured.NestedString(dsc.Object, "status", "components", component, "managementState")
+	if err != nil || !found {
+		return ""
+	}
+	return state
+}
+
+func ComponentConditionStatus(dsc *unstructured.Unstructured, conditionType string) string {
+	conditions, found, err := unstructured.NestedSlice(dsc.Object, "status", "conditions")
+	if err != nil || !found {
+		return ""
+	}
+	for _, c := range conditions {
+		condition, ok := c.(map[string]interface{})
+		if !ok {
+			continue
+		}
+		if condition["type"] == conditionType {
+			if status, ok := condition["status"].(string); ok {
+				return status
+			}
+		}
+	}
+	return ""
+}

tests/common/support/dataScienceCluster.go

@@ -129,3 +129,21 @@ func KueueWorkloadAdmitted(workload *kueuev1beta1.Workload) bool {
 	}
 	return false
 }
+
+func KueueWorkloadEvicted(workload *kueuev1beta1.Workload) bool {
+	for _, v := range workload.Status.Conditions {
+		if v.Type == "Evicted" && v.Status == "True" {
+			return true
+		}
+	}
+	return false
+}
+
+func KueueWorkloadInadmissible(workload *kueuev1beta1.Workload) (bool, string) {
+	for _, v := range workload.Status.Conditions {
+		if v.Type == "QuotaReserved" && v.Status == "False" && v.Reason == "Inadmissible" {
+			return true, v.Message
+		}
+	}
+	return false, ""
+}