[RayCluster][Fix] Add expectations of RayCluster

ray-operator/controllers/ray/expectations/scale_expectation.go

@@ -0,0 +1,152 @@
+package expectations
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/tools/cache"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+const DefaultHead = ""
+
+// ScaleAction is the action of scale, like create and delete.
+type ScaleAction string
+
+const (
+	// Create action
+	Create ScaleAction = "create"
+	// Delete action
+	Delete ScaleAction = "delete"
+)
+
+const (
+	// GroupIndex indexes pods within the specified group in RayCluster
+	GroupIndex = "group"
+	// RayClusterIndex indexes pods within the RayCluster
+	RayClusterIndex = "raycluster"
+)
+
+// RayClusterScaleExpectation is an interface that to set and wait on expectations of RayCluster groups scale.
+type RayClusterScaleExpectation interface {
+	ExpectScalePod(rayClusterName, group, namespace, name string, action ScaleAction)
+	IsSatisfied(ctx context.Context, rayClusterName, group, namespace string) bool
+	Delete(rayClusterName, namespace string)
+}
+
+func newRayPodIndexer() cache.Indexer {
+	return cache.NewIndexer(rayPodKey, cache.Indexers{GroupIndex: groupIndexFunc, RayClusterIndex: rayClusterIndexFunc})
+}
+
+func NewRayClusterScaleExpectation(client client.Client) RayClusterScaleExpectation {
+	return &realRayClusterScaleExpectation{
+		Client:     client,
+		itemsCache: newRayPodIndexer(),
+	}
+}
+
+type realRayClusterScaleExpectation struct {
+	client.Client
+	itemsCache cache.Indexer
+}
+
+func (r *realRayClusterScaleExpectation) ExpectScalePod(rayClusterName, group, namespace, name string, action ScaleAction) {
+	_ = r.itemsCache.Add(&rayPod{
+		name:            name,
+		namespace:       namespace,
+		group:           group,
+		rayCluster:      rayClusterName,
+		action:          action,
+		recordTimestamp: time.Now(),
+	})
+}
+
+func (r *realRayClusterScaleExpectation) IsSatisfied(ctx context.Context, rayClusterName, group, namespace string) (isSatisfied bool) {
+	items, _ := r.itemsCache.ByIndex(GroupIndex, fmt.Sprintf("%s/%s/%s", namespace, rayClusterName, group))
+	isSatisfied = true
+	for i := range items {
+		rp := items[i].(*rayPod)
+		pod := &corev1.Pod{}
+		isPodSatisfied := false
+		switch rp.action {
+		case Create:
+			if err := r.Get(ctx, types.NamespacedName{Name: rp.name, Namespace: namespace}, pod); err == nil {
+				isPodSatisfied = true
+			} else {
+				isPodSatisfied = errors.IsNotFound(err) && rp.recordTimestamp.Add(30*time.Second).Before(time.Now())
+			}
+		case Delete:
+			if err := r.Get(ctx, types.NamespacedName{Name: rp.name, Namespace: namespace}, pod); err != nil {
+				isPodSatisfied = errors.IsNotFound(err)
+			} else {
+				isPodSatisfied = pod.DeletionTimestamp != nil
+			}
+		}
+		// delete satisfied item in cache
+		if isPodSatisfied {
+			_ = r.itemsCache.Delete(items[i])
+		} else {
+			isSatisfied = false
+		}
+	}
+	return isSatisfied
+}
+
+func (r *realRayClusterScaleExpectation) Delete(rayClusterName, namespace string) {
+	items, _ := r.itemsCache.ByIndex(RayClusterIndex, fmt.Sprintf("%s/%s", namespace, rayClusterName))
+	for _, item := range items {
+		_ = r.itemsCache.Delete(item)
+	}
+}
+
+type rayPod struct {
+	recordTimestamp time.Time
+	action          ScaleAction
+	name            string
+	namespace       string
+	rayCluster      string
+	group           string
+}
+
+func (p *rayPod) Key() string {
+	return fmt.Sprintf("%s/%s", p.namespace, p.name)
+}
+
+func (p *rayPod) GroupKey() string {
+	return fmt.Sprintf("%s/%s/%s", p.namespace, p.rayCluster, p.group)
+}
+
+func (p *rayPod) ClusterKey() string {
+	return fmt.Sprintf("%s/%s", p.namespace, p.rayCluster)
+}
+
+func rayPodKey(obj interface{}) (string, error) {
+	return obj.(*rayPod).Key(), nil
+}
+
+func groupIndexFunc(obj interface{}) ([]string, error) {
+	return []string{obj.(*rayPod).GroupKey()}, nil
+}
+
+func rayClusterIndexFunc(obj interface{}) ([]string, error) {
+	return []string{obj.(*rayPod).ClusterKey()}, nil
+}
+
+func NewFakeRayClusterScaleExpectation() RayClusterScaleExpectation {
+	return &fakeRayClusterScaleExpectation{}
+}
+
+type fakeRayClusterScaleExpectation struct{}
+
+func (r *fakeRayClusterScaleExpectation) ExpectScalePod(_, _, _, _ string, _ ScaleAction) {
+}
+
+func (r *fakeRayClusterScaleExpectation) IsSatisfied(_ context.Context, _, _, _ string) bool {
+	return true
+}
+
+func (r *fakeRayClusterScaleExpectation) Delete(_, _ string) {}

ray-operator/controllers/ray/expectations/scale_expectation_test.go

@@ -0,0 +1,87 @@
+package expectations
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	clientFake "sigs.k8s.io/controller-runtime/pkg/client/fake"
+)
+
+func TestRayClusterExpectations(t *testing.T) {
+	setupTest()
+	ctx := context.TODO()
+	fakeClient := clientFake.NewClientBuilder().WithRuntimeObjects().Build()
+	exp := NewRayClusterScaleExpectation(fakeClient)
+	namespace := "default"
+	rayClusterName := "raycluster-test-pod"
+
+	// Test expect create head
+	exp.ExpectScalePod(rayClusterName, DefaultHead, namespace, testPods[0].(*corev1.Pod).Name, Create)
+	assert.Equal(t, len(exp.(*realRayClusterScaleExpectation).itemsCache.List()), 1)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), false)
+	err := fakeClient.Create(context.TODO(), testPods[0].(*corev1.Pod))
+	assert.Nil(t, err, "Fail to create head pod")
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), true)
+	// delete satisfied item in cache
+	assert.Equal(t, len(exp.(*realRayClusterScaleExpectation).itemsCache.List()), 0)
+
+	// Test expect delete head
+	exp.ExpectScalePod(rayClusterName, DefaultHead, namespace, testPods[0].(*corev1.Pod).Name, Delete)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), false)
+	// delete pod
+	err = fakeClient.Delete(context.TODO(), testPods[0].(*corev1.Pod))
+	assert.Nil(t, err, "Fail to delete head pod")
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), true)
+
+	// Test expect create worker
+	group := "test-group"
+	exp.ExpectScalePod(rayClusterName, group, namespace, testPods[1].(*corev1.Pod).Name, Create)
+	exp.ExpectScalePod(rayClusterName, group, namespace, testPods[2].(*corev1.Pod).Name, Create)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, group, namespace), false)
+	assert.Nil(t, fakeClient.Create(context.TODO(), testPods[1].(*corev1.Pod)), "Fail to create worker pod1")
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, group, namespace), false)
+	assert.Nil(t, fakeClient.Create(context.TODO(), testPods[2].(*corev1.Pod)), "Fail to create worker pod2")
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, group, namespace), true)
+
+	// Test delete all
+	// reset pods
+	setupTest()
+	exp.ExpectScalePod(rayClusterName, DefaultHead, namespace, testPods[0].(*corev1.Pod).Name, Create)
+	exp.ExpectScalePod(rayClusterName, group, namespace, testPods[1].(*corev1.Pod).Name, Delete)
+	exp.ExpectScalePod(rayClusterName, group, namespace, testPods[2].(*corev1.Pod).Name, Delete)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), false)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, group, namespace), false)
+	exp.Delete(rayClusterName, namespace)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, DefaultHead, namespace), true)
+	assert.Equal(t, exp.IsSatisfied(ctx, rayClusterName, group, namespace), true)
+	assert.Equal(t, len(exp.(*realRayClusterScaleExpectation).itemsCache.List()), 0)
+}
+
+var testPods []runtime.Object
+
+func setupTest() {
+	testPods = []runtime.Object{
+		&corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "pod1",
+				Namespace: "default",
+			},
+		},
+		&corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "pod2",
+				Namespace: "default",
+			},
+		},
+		&corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "pod3",
+				Namespace: "default",
+			},
+		},
+	}
+}