wip env manager with kafka

Author: MiguelAAe

tests/integration/godog/k8sclient/component_manager.go

@@ -0,0 +1,60 @@
+package k8sclient
+
+import (
+	"context"
+	"sync"
+)
+
+type ComponentManager interface {
+	Name() ComponentName
+	Snapshot(ctx context.Context) error
+	Restore(ctx context.Context) error
+}
+
+type EnvManager struct {
+	mu         sync.Mutex
+	components map[ComponentName]ComponentManager
+}
+
+type ComponentName string
+
+func NewEnvManager(components ...ComponentManager) *EnvManager {
+	comps := make(map[ComponentName]ComponentManager, len(components))
+	for _, component := range components {
+		comps[component.Name()] = component
+	}
+	return &EnvManager{components: comps}
+}
+
+func (e *EnvManager) SnapshotAll(ctx context.Context) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	for _, c := range e.components {
+		if err := c.Snapshot(ctx); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (e *EnvManager) RestoreAll(ctx context.Context) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	for _, c := range e.components {
+		if err := c.Restore(ctx); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (e *EnvManager) Kafka() *KafkaComponent {
+	for _, component := range e.components {
+		if k, ok := component.(*KafkaComponent); ok {
+			return k
+		}
+	}
+	return nil
+}

tests/integration/godog/k8sclient/kafka_component.go

@@ -0,0 +1,102 @@
+package k8sclient
+
+import (
+	"context"
+	"fmt"
+
+	appsv1 "k8s.io/api/apps/v1"
+	"k8s.io/apimachinery/pkg/types"
+)
+
+type KafkaComponent struct {
+	name      ComponentName
+	namespace string
+	resource  types.NamespacedName
+
+	k8s   *K8sClient
+	base  *int32 // baseline replicas
+	dirty bool
+}
+
+const kafkaComponentName ComponentName = "kafka"
+
+func NewKafkaComponent(k8s *K8sClient, namespace, statefulSetName string) *KafkaComponent {
+	return &KafkaComponent{
+		name:      kafkaComponentName,
+		namespace: namespace,
+		resource:  types.NamespacedName{Namespace: namespace, Name: statefulSetName},
+		k8s:       k8s,
+	}
+}
+
+func (k *KafkaComponent) Name() ComponentName { return k.name }
+
+// Snapshot reads the current replicas once and stores them as baseline.
+func (k *KafkaComponent) Snapshot(ctx context.Context) error {
+	if k.base != nil {
+		// already snapshotted
+		return nil
+	}
+
+	var sts appsv1.StatefulSet
+	if err := k.k8s.KubeClient.Get(ctx, k.resource, &sts); err != nil {
+		return fmt.Errorf("kafka snapshot: get %s/%s: %w", k.resource.Namespace, k.resource.Name, err)
+	}
+
+	if sts.Spec.Replicas == nil {
+		return fmt.Errorf("kafka snapshot: statefulset %s/%s has nil .spec.replicas", k.resource.Namespace, k.resource.Name)
+	}
+
+	replicas := *sts.Spec.Replicas
+	k.base = &replicas
+	return nil
+}
+
+// Restore only acts if we've modified Kafka in this scenario (dirty=true).
+func (k *KafkaComponent) Restore(ctx context.Context) error {
+	if !k.dirty || k.base == nil {
+		return nil
+	}
+	if err := k.scale(ctx, *k.base); err != nil {
+		return err
+	}
+	k.dirty = false
+	return nil
+}
+
+// Public helpers used from steps:
+
+func (k *KafkaComponent) MakeUnavailable(ctx context.Context) error {
+	if err := k.Snapshot(ctx); err != nil {
+		return err
+	}
+	if err := k.scale(ctx, 0); err != nil {
+		return err
+	}
+	k.dirty = true
+	return nil
+}
+
+func (k *KafkaComponent) MakeAvailable(ctx context.Context) error {
+	if k.base == nil {
+		// nothing to restore to; safest is to no-op or return error
+		return nil
+	}
+	if err := k.scale(ctx, *k.base); err != nil {
+		return err
+	}
+	k.dirty = false
+	return nil
+}
+
+func (k *KafkaComponent) scale(ctx context.Context, replicas int32) error {
+	var sts appsv1.StatefulSet
+	if err := k.k8s.KubeClient.Get(ctx, k.resource, &sts); err != nil {
+		return fmt.Errorf("kafka scale: get %s/%s: %w", k.resource.Namespace, k.resource.Name, err)
+	}
+	sts.Spec.Replicas = &replicas
+	if err := k.k8s.KubeClient.Update(ctx, &sts); err != nil {
+		return fmt.Errorf("kafka scale: update %s/%s: %w", k.resource.Namespace, k.resource.Name, err)
+	}
+	return nil
+}

tests/integration/godog/steps/env_steps.go

@@ -0,0 +1,47 @@
+package steps
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cucumber/godog"
+	"github.com/seldonio/seldon-core/tests/integration/godog/k8sclient"
+)
+
+type Env struct {
+	env *k8sclient.EnvManager
+}
+
+func LoadEnvSteps(scenario *godog.ScenarioContext, w *World) {
+	scenario.Step(`^(Kafka) is unavailable for Core 2 with timeout "([^"]+)"`, func(kind, timeout string) error {
+		return withTimeoutCtx(timeout, func(ctx context.Context) error {
+			switch kind {
+			case "Kafka":
+				k := w.env.Kafka()
+				if k != nil {
+					return nil
+				}
+
+				return k.MakeUnavailable(ctx)
+			default:
+				return fmt.Errorf("unknown target type: %s", kind)
+			}
+		})
+	})
+
+	scenario.Step(`^(Kafka) is available for Core 2 with timeout "([^"]+)"`, func(kind, timeout string) error {
+		return withTimeoutCtx(timeout, func(ctx context.Context) error {
+			switch kind {
+			case "Kafka":
+				k := w.env.Kafka()
+				if k != nil {
+					return nil
+				}
+
+				return k.MakeAvailable(ctx)
+			default:
+				return fmt.Errorf("unknown target type: %s", kind)
+			}
+		})
+	})
+}

tests/integration/godog/steps/world.go

@@ -19,13 +19,11 @@ import (
 )
 
 type World struct {
-	namespace            string
-	kubeClient           *k8sclient.K8sClient
-	corek8sClient        v.Interface
-	watcherStorage       k8sclient.WatcherStorage
-	StartingClusterState string //todo: this will be a combination of starting state awareness of core 2 such as the
-	//todo:  server config,seldon config and seldon runtime to be able to reconcile to starting state should we change
-	//todo: the state such as reducing replicas to 0 of scheduler to test unavailability
+	namespace         string
+	kubeClient        *k8sclient.K8sClient
+	corek8sClient     v.Interface
+	watcherStorage    k8sclient.WatcherStorage
+	env               *k8sclient.EnvManager //this is a combination of components for the cluster
 	currentModel      *Model
 	currentPipeline   *Pipeline
 	currentExperiment *Experiment
@@ -41,6 +39,7 @@ type Config struct {
 	KubeClient     *k8sclient.K8sClient
 	K8sClient      v.Interface
 	WatcherStorage k8sclient.WatcherStorage
+	Env            *k8sclient.EnvManager
 	GRPC           v2_dataplane.GRPCInferenceServiceClient
 	IngressHost    string
 	HTTPPort       uint
@@ -62,6 +61,7 @@ func NewWorld(c Config) (*World, error) {
 		namespace:         c.Namespace,
 		kubeClient:        c.KubeClient,
 		watcherStorage:    c.WatcherStorage,
+		env:               c.Env,
 		currentModel:      newModel(label, c.Namespace, c.K8sClient, c.Logger, c.WatcherStorage),
 		currentExperiment: newExperiment(label, c.Namespace, c.K8sClient, c.Logger, c.WatcherStorage),
 		currentPipeline:   newPipeline(label, c.Namespace, c.K8sClient, c.Logger, c.WatcherStorage),

tests/integration/godog/suite/suite.go

@@ -13,6 +13,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"time"
 
 	"github.com/cucumber/godog"
 	v2dataplane "github.com/seldonio/seldon-core/apis/go/v2/mlops/v2_dataplane"
@@ -32,6 +33,7 @@ type dependencies struct {
 	mlopsClient     *v.Clientset
 	watcherStore    k8sclient.WatcherStorage
 	inferenceClient v2dataplane.GRPCInferenceServiceClient
+	env             *k8sclient.EnvManager
 	Config          *GodogConfig
 }
 
@@ -53,6 +55,8 @@ type dependencies struct {
 var suiteDeps dependencies
 
 func InitializeTestSuite(ctx *godog.TestSuiteContext) {
+	realContext, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
 	// Load configuration from JSON file
 	config, err := LoadConfig()
 	if err != nil {
@@ -101,12 +105,23 @@ func InitializeTestSuite(ctx *godog.TestSuiteContext) {
 	}
 	grpcClient := v2dataplane.NewGRPCInferenceServiceClient(conn)
 
+	// Example: kafka in namespace "mlops" with statefulset name "kafka"
+	kafka := k8sclient.NewKafkaComponent(k8sClient, config.Namespace, "kafka")
+
+	env := k8sclient.NewEnvManager(kafka)
+
+	// Snapshot baseline state once per suite
+	if err := env.SnapshotAll(realContext); err != nil {
+		panic(fmt.Errorf("failed to snapshot environment: %w", err))
+	}
+
 	suiteDeps.logger = log
 	suiteDeps.k8sClient = k8sClient
 	suiteDeps.mlopsClient = clientSet // todo: this clientSet might get use for get requests or for the mlops interface and could be passed to the world might be split up by type
 	suiteDeps.watcherStore = watchStore
 	suiteDeps.Config = config
 	suiteDeps.inferenceClient = grpcClient
+	suiteDeps.env = env
 
 	ctx.BeforeSuite(func() {
 		suiteDeps.watcherStore.Start()
@@ -118,6 +133,11 @@ func InitializeTestSuite(ctx *godog.TestSuiteContext) {
 
 	ctx.AfterSuite(func() {
 		suiteDeps.watcherStore.Stop()
+
+		err := suiteDeps.env.RestoreAll(context.Background())
+		if err != nil {
+			return
+		}
 		// e.g. clean namespace, close clients if needed delete servers
 	})
 }
@@ -172,5 +192,5 @@ func InitializeScenario(scenarioCtx *godog.ScenarioContext) {
 	steps.LoadServerSteps(scenarioCtx, world)
 	steps.LoadCustomPipelineSteps(scenarioCtx, world)
 	steps.LoadExperimentSteps(scenarioCtx, world)
-	// TODO: load other steps, e.g. pipeline, experiment, etc.
+	steps.LoadEnvSteps(scenarioCtx, world)
 }