Python runtime sandbox e2e test

test/e2e/framework/client.go

@@ -29,11 +29,18 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/agent-sandbox/test/e2e/framework/predicates"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
+const (
+	// DefaultTimeout is the default timeout for WaitForObject.
+	DefaultTimeout = 60 * time.Second
+)
+
 // ClusterClient is an abstraction layer for test cases to interact with the cluster.
 type ClusterClient struct {
 	*testing.T
@@ -111,9 +118,11 @@ func (cl *ClusterClient) ValidateObjectNotFound(ctx context.Context, obj client.
 // predicates.
 func (cl *ClusterClient) WaitForObject(ctx context.Context, obj client.Object, p ...predicates.ObjectPredicate) error {
 	cl.Helper()
-	// Static 30 second timeout, this can be adjusted if needed
-	timeoutCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
+	var cancel context.CancelFunc
+	if _, ok := ctx.Deadline(); !ok {
+		ctx, cancel = context.WithTimeout(ctx, DefaultTimeout)
+		defer cancel()
+	}
 	start := time.Now()
 	nn := types.NamespacedName{Name: obj.GetName(), Namespace: obj.GetNamespace()}
 	defer func() {
@@ -123,10 +132,10 @@ func (cl *ClusterClient) WaitForObject(ctx context.Context, obj client.Object, p
 	var validationErr error
 	for {
 		select {
-		case <-timeoutCtx.Done():
+		case <-ctx.Done():
 			return fmt.Errorf("timed out waiting for object: %w", validationErr)
 		default:
-			if validationErr = cl.ValidateObject(timeoutCtx, obj, p...); validationErr == nil {
+			if validationErr = cl.ValidateObject(ctx, obj, p...); validationErr == nil {
 				return nil
 			}
 			// Simple sleep for fixed duration (basic MVP)
@@ -253,3 +262,62 @@ func (cl *ClusterClient) PortForward(ctx context.Context, pod types.NamespacedNa
 	}
 	return nil
 }
+
+var sandboxGVK = schema.GroupVersionKind{
+	Group:   "agents.x-k8s.io",
+	Version: "v1alpha1",
+	Kind:    "Sandbox",
+}
+
+var sandboxWarmpoolGVK = schema.GroupVersionKind{
+	Group:   "extensions.agents.x-k8s.io",
+	Version: "v1alpha1",
+	Kind:    "SandboxWarmPool",
+}
+
+func (cl *ClusterClient) WaitForSandboxReady(ctx context.Context, sandboxID types.NamespacedName) error {
+	sandbox := &unstructured.Unstructured{}
+	sandbox.SetGroupVersionKind(sandboxGVK)
+	sandbox.SetName(sandboxID.Name)
+	sandbox.SetNamespace(sandboxID.Namespace)
+	timeoutCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
+	defer cancel()
+	if err := cl.WaitForObject(timeoutCtx, sandbox, predicates.ReadyConditionIsTrue); err != nil {
+		cl.Logf("waiting for sandbox to be ready: %v", err)
+		return err
+	}
+	return nil
+}
+
+func (cl *ClusterClient) WaitForWarmPoolReady(ctx context.Context, sandboxWarmpoolID types.NamespacedName, expectedReplicas int) error {
+	cl.Helper()
+	cl.Logf("Waiting for SandboxWarmPool Pods to be ready: warmpoolID - %s; expectedReplicas - %d", sandboxWarmpoolID, expectedReplicas)
+
+	warmpool := &unstructured.Unstructured{}
+	warmpool.SetGroupVersionKind(sandboxWarmpoolGVK)
+	if err := cl.client.Get(ctx, sandboxWarmpoolID, warmpool); err != nil {
+		cl.T.Fatalf("Failed to get SandboxWarmPool %s: %v", sandboxWarmpoolID, err)
+		return err
+	}
+
+	if err := cl.WaitForObject(ctx, warmpool, predicates.ReadyReplicasConditionIsTrue); err != nil {
+		cl.T.Fatalf("waiting for warmpool to be ready: %v", err)
+		return err
+	}
+	return nil
+
+}
+
+// GetSandbox returns the Sandbox object from the cluster.
+func (cl *ClusterClient) GetSandbox(ctx context.Context, sandboxID types.NamespacedName) (*unstructured.Unstructured, error) {
+	sandbox := &unstructured.Unstructured{}
+	sandbox.SetGroupVersionKind(sandboxGVK)
+	sandbox.SetName(sandboxID.Name)
+	sandbox.SetNamespace(sandboxID.Namespace)
+
+	if err := cl.client.Get(ctx, sandboxID, sandbox); err != nil {
+		cl.Logf("failed to get Sandbox %s: %v", sandboxID, err)
+		return nil, err
+	}
+	return sandbox, nil
+}

test/e2e/framework/context.go

@@ -23,6 +23,7 @@ import (
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/client-go/tools/clientcmd"
 	"sigs.k8s.io/agent-sandbox/controllers"
+	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -37,6 +38,7 @@ var (
 
 func init() {
 	utilruntime.Must(apiextensionsv1.AddToScheme(controllers.Scheme))
+	utilruntime.Must(extensionsv1alpha1.AddToScheme(controllers.Scheme))
 }
 
 // TestContext is a helper for managing e2e test scaffolding.

test/e2e/framework/predicates/conditions.go

@@ -26,7 +26,8 @@ import (
 // objectWithStatus is a simplified struct to parse the status of a resource.
 type objectWithStatus struct {
 	Status struct {
-		Conditions []metav1.Condition `json:"conditions,omitempty"`
+		Conditions    []metav1.Condition `json:"conditions,omitempty"`
+		ReadyReplicas int                `json:"readyReplicas,omitempty"`
 	} `json:"status"`
 }
 
@@ -62,3 +63,20 @@ func asUnstructured(obj client.Object) (*unstructured.Unstructured, error) {
 	}
 	return &unstructured.Unstructured{Object: m}, nil
 }
+
+// ReadyReplicasConditionIsTrue checks if the given object has more than 0 replicas.
+func ReadyReplicasConditionIsTrue(obj client.Object) error {
+	u, err := asUnstructured(obj)
+	if err != nil {
+		return fmt.Errorf("failed to convert to unstructured: %w", err)
+	}
+
+	var status objectWithStatus
+	if err := runtime.DefaultUnstructuredConverter.FromUnstructured(u.Object, &status); err != nil {
+		return fmt.Errorf("failed to convert to objectWithStatus: %v", err)
+	}
+	if status.Status.ReadyReplicas > 0 {
+		return nil
+	}
+	return fmt.Errorf("object does not have more than 0 replicas: %d", status.Status.ReadyReplicas)
+}