Merge remote-tracking branch 'upstream/main'

Author: m-rafeeq

tests/common/support/dataScienceCluster.go

@@ -19,16 +19,36 @@ package support
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/onsi/gomega"
 
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	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"
 )
 
+// DSC constants
+const (
+	// Phase values
+	PhaseReady = "Ready"
+
+	// Condition status values
+	ConditionTrue  = "True"
+	ConditionFalse = "False"
+
+	// Condition reason values
+	ReasonRemoved = "Removed"
+
+	// Management state values
+	StateManaged   = "Managed"
+	StateRemoved   = "Removed"
+	StateUnmanaged = "Unmanaged"
+)
+
 var DscGVR = schema.GroupVersionResource{
 	Group:    "datasciencecluster.opendatahub.io",
 	Version:  "v2",
@@ -43,50 +63,187 @@ func updateDSC(dynamicClient dynamic.Interface, ctx context.Context, dsc *unstru
 	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
+func GetDSCPhase(dsc *unstructured.Unstructured) string {
+	phase, found, err := unstructured.NestedString(dsc.Object, "status", "phase")
+	if err != nil || !found {
+		return ""
 	}
+	return phase
+}
 
-	err = unstructured.SetNestedField(dsc.Object, state, "spec", "components", component, "managementState")
-	if err != nil {
-		return err
+func WaitForDSCReady(dynamicClient dynamic.Interface, ctx context.Context, dscName string) error {
+	var phase string
+
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("%s: timed out waiting for DSC to be Ready - current state: %s", dscName, phase)
+		default:
+			dsc, err := getDSC(dynamicClient, ctx, dscName)
+			if err != nil {
+				return err
+			}
+			phase = GetDSCPhase(dsc)
+			if phase == PhaseReady {
+				return nil
+			}
+			time.Sleep(5 * time.Second)
+		}
+	}
+}
+
+type ConditionDetails struct {
+	Type    string
+	Status  string
+	Reason  string
+	Message string
+}
+
+func componentReadyCondition(component string) string {
+	if len(component) == 0 {
+		return ""
 	}
+	return strings.ToUpper(component[:1]) + component[1:] + "Ready"
+}
 
-	_, err = updateDSC(dynamicClient, ctx, dsc)
-	return err
+func GetConditionDetails(dsc *unstructured.Unstructured, conditionType string) ConditionDetails {
+	details := ConditionDetails{Type: conditionType}
+	conditions, found, err := unstructured.NestedSlice(dsc.Object, "status", "conditions")
+	if err != nil || !found {
+		return details
+	}
+	for _, c := range conditions {
+		condition, ok := c.(map[string]interface{})
+		if !ok {
+			continue
+		}
+		if condition["type"] == conditionType {
+			if status, ok := condition["status"].(string); ok {
+				details.Status = status
+			}
+			if reason, ok := condition["reason"].(string); ok {
+				details.Reason = reason
+			}
+			if message, ok := condition["message"].(string); ok {
+				details.Message = message
+			}
+			return details
+		}
+	}
+	return details
 }
 
-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) {
+func setComponentState(dynamicClient dynamic.Interface, ctx context.Context, dscName, component, state string) error {
+	// Retry up to 5 times to handle optimistic locking conflicts
+	maxRetries := 5
+	for i := 0; i < maxRetries; i++ {
 		dsc, err := getDSC(dynamicClient, ctx, dscName)
 		if err != nil {
 			return err
 		}
-		currentState := ComponentStatusManagementState(dsc, component)
-		if currentState == expectedState {
+
+		err = unstructured.SetNestedField(dsc.Object, state, "spec", "components", component, "managementState")
+		if err != nil {
+			return err
+		}
+
+		_, err = updateDSC(dynamicClient, ctx, dsc)
+		if err == nil {
 			return nil
 		}
-		time.Sleep(2 * time.Second)
+
+		if apierrors.IsConflict(err) {
+			time.Sleep(500 * time.Millisecond)
+			continue
+		}
+
+		return err
+	}
+	return fmt.Errorf("failed to set component %s state after %d retries due to conflicts", component, maxRetries)
+}
+
+func WaitForComponentReady(dynamicClient dynamic.Interface, ctx context.Context, dscName, component string) error {
+	conditionType := componentReadyCondition(component)
+	var condition ConditionDetails
+
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("%s: timed out waiting for component to be ready - status: %s, reason: %s, message: %s",
+				conditionType, condition.Status, condition.Reason, condition.Message)
+		default:
+			dsc, err := getDSC(dynamicClient, ctx, dscName)
+			if err != nil {
+				return err
+			}
+			condition = GetConditionDetails(dsc, conditionType)
+
+			if condition.Status == ConditionTrue {
+				return nil
+			}
+
+			time.Sleep(5 * time.Second)
+		}
 	}
-	return fmt.Errorf("timeout waiting for component %s to reach state %s", component, expectedState)
 }
 
+func WaitForComponentRemoved(dynamicClient dynamic.Interface, ctx context.Context, dscName, component string) error {
+	conditionType := componentReadyCondition(component)
+	var condition ConditionDetails
+
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("%s: timed out waiting for component to be removed - status: %s, reason: %s, message: %s",
+				conditionType, condition.Status, condition.Reason, condition.Message)
+		default:
+			dsc, err := getDSC(dynamicClient, ctx, dscName)
+			if err != nil {
+				return err
+			}
+			condition = GetConditionDetails(dsc, conditionType)
+
+			if condition.Status == "False" && condition.Reason == ReasonRemoved {
+				return nil
+			}
+
+			time.Sleep(5 * time.Second)
+		}
+	}
+}
+
+// SetComponentStateAndWait sets a component's managementState and waits for completion.
+// It creates a context with timeout that enforces the deadline across all operations.
 func SetComponentStateAndWait(dynamicClient dynamic.Interface, ctx context.Context, dscName, component, state string, timeout time.Duration) error {
+	// Create a context with timeout - this will automatically enforce the deadline
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	if err := setComponentState(dynamicClient, ctx, dscName, component, state); err != nil {
 		return err
 	}
-	return WaitForComponentState(dynamicClient, ctx, dscName, component, state, timeout)
+
+	if state == StateManaged {
+		if err := WaitForComponentReady(dynamicClient, ctx, dscName, component); err != nil {
+			return err
+		}
+	}
+
+	if state == StateRemoved {
+		if err := WaitForComponentRemoved(dynamicClient, ctx, dscName, component); err != nil {
+			return err
+		}
+	}
+
+	return WaitForDSCReady(dynamicClient, ctx, dscName)
 }
 
 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")
+	return setComponentState(test.Client().Dynamic(), test.Ctx(), dscName, component, StateUnmanaged)
 }
 
 func DSCResource(test Test, name string) func(g gomega.Gomega) *unstructured.Unstructured {

tests/trainer/trainer_kueue_integration_test.go

@@ -46,6 +46,9 @@ var initialKueueState string
 var initialTrainerState string
 
 func TestMain(m *testing.M) {
+	var code int
+	var setupFailed bool
+
 	// Capture initial Trainer state before running any tests
 	initialTrainerState = captureComponentState("trainer")
 	fmt.Printf("Initial Trainer managementState: %s\n", initialTrainerState)
@@ -57,15 +60,19 @@ func TestMain(m *testing.M) {
 	// Setup Trainer to Managed if not already
 	if initialTrainerState != "Managed" {
 		if err := setupTrainer(); err != nil {
-			fmt.Printf("Setup: Failed to set Trainer managementState to Managed in DataScienceCluster: %v\n", err)
-			os.Exit(1)
+			fmt.Printf("Setup failed: %v\n", err)
+			fmt.Println("Skipping test execution due to setup failure ...")
+			setupFailed = true
+			code = 1
 		}
 	} else {
 		fmt.Println("Setup: Skipping Trainer setup as it is already set to Managed in DataScienceCluster")
 	}
 
-	// Run all tests
-	code := m.Run()
+	// Run all tests only if setup succeeded
+	if !setupFailed {
+		code = m.Run()
+	}
 
 	// TearDown Trainer: Only set to Removed if it was not already Managed before tests
 	if initialTrainerState != "Managed" {
@@ -79,7 +86,7 @@ func TestMain(m *testing.M) {
 	// TearDown Kueue: Only set to Removed if it was not already Unmanaged before tests
 	if initialKueueState != "Unmanaged" {
 		if err := tearDownComponent("kueue"); err != nil {
-			fmt.Printf("TearDown: Failed to set Kueue to Removed : %v\n", err)
+			fmt.Printf("TearDown: Failed to set Kueue to Removed: %v\n", err)
 		}
 	} else {
 		fmt.Println("TearDown: Skipping Kueue teardown as Initial Kueue managementState was Unmanaged in DataScienceCluster")
@@ -124,13 +131,13 @@ func captureComponentState(component string) string {
 func setupTrainer() error {
 	dynamicClient, err := createDynamicClient()
 	if err != nil {
-		return fmt.Errorf("Setup: %w", err)
+		return err
 	}
 
 	fmt.Println("Setup: Setting trainer managementState to Managed in DataScienceCluster...")
 	err = SetComponentStateAndWait(dynamicClient, context.Background(), defaultDSCName, "trainer", "Managed", 2*time.Minute)
 	if err != nil {
-		return fmt.Errorf("Setup: failed to set trainer to Managed: %w", err)
+		return err
 	}
 
 	fmt.Println("Setup: Trainer is set to Managed managementState successfully")