From e15a9fed762e06cc59b0fa340d45c4f564c70499 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Thu, 27 Jun 2024 11:24:51 +0300 Subject: [PATCH 1/6] [jobframework] Only check if the owner is managed for enabled integrations. --- .../jobframework/integrationmanager.go | 32 +++++++++++++++++-- .../jobframework/integrationmanager_test.go | 16 +++++++++- pkg/controller/jobframework/setup.go | 1 + .../jobs/job/job_controller_test.go | 1 + .../jobs/pod/pod_controller_test.go | 1 + pkg/controller/jobs/pod/pod_webhook_test.go | 28 ++++++++++++++++ .../controller/jobs/job/suite_test.go | 2 ++ .../controller/jobs/mpijob/suite_test.go | 1 + .../controller/jobs/pod/suite_test.go | 1 + .../controller/jobs/raycluster/suite_test.go | 2 ++ 10 files changed, 82 insertions(+), 3 deletions(-) diff --git a/pkg/controller/jobframework/integrationmanager.go b/pkg/controller/jobframework/integrationmanager.go index 894f8f5ddd..2b0cf61da4 100644 --- a/pkg/controller/jobframework/integrationmanager.go +++ b/pkg/controller/jobframework/integrationmanager.go @@ -26,6 +26,7 @@ import ( "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/client-go/tools/record" + "k8s.io/utils/set" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/reconcile" @@ -72,6 +73,7 @@ type IntegrationCallbacks struct { type integrationManager struct { names []string integrations map[string]IntegrationCallbacks + enabledIntegrations set.Set[string] externalIntegrations map[string]runtime.Object } @@ -144,6 +146,14 @@ func (m *integrationManager) getExternal(kindArg string) (runtime.Object, bool) return jt, f } +func (m *integrationManager) enableIntegration(name string) { + if m.enabledIntegrations == nil { + m.enabledIntegrations = set.New(name) + } else { + m.enabledIntegrations.Insert(name) + } +} + func (m *integrationManager) getList() []string { ret := make([]string, len(m.names)) copy(ret, m.names) @@ -152,8 +162,9 @@ func (m *integrationManager) getList() []string { } func (m *integrationManager) getJobTypeForOwner(ownerRef *metav1.OwnerReference) runtime.Object { - for _, cbs := range m.integrations { - if cbs.IsManagingObjectsOwner != nil && cbs.IsManagingObjectsOwner(ownerRef) { + for jobKey := range m.enabledIntegrations { + cbs, found := m.integrations[jobKey] + if found && cbs.IsManagingObjectsOwner != nil && cbs.IsManagingObjectsOwner(ownerRef) { return cbs.JobType } } @@ -186,6 +197,23 @@ func ForEachIntegration(f func(name string, cb IntegrationCallbacks) error) erro return manager.forEach(f) } +// EnableIntegration marks the integration identified by name as enabled. +func EnableIntegration(name string) { + manager.enableIntegration(name) +} + +// EnableIntegrationsForTest - should be used only in tests +// Mark the frameworks identified by names and return a revert function. +func EnableIntegrationsForTest(names ...string) func() { + old := manager.enabledIntegrations.Clone() + for _, name := range names { + manager.enableIntegration(name) + } + return func() { + manager.enabledIntegrations = old + } +} + // GetIntegration looks-up the framework identified by name in the currently registered // list of frameworks returning its callbacks and true if found. func GetIntegration(name string) (IntegrationCallbacks, bool) { diff --git a/pkg/controller/jobframework/integrationmanager_test.go b/pkg/controller/jobframework/integrationmanager_test.go index c3f124599b..d6777a4b6c 100644 --- a/pkg/controller/jobframework/integrationmanager_test.go +++ b/pkg/controller/jobframework/integrationmanager_test.go @@ -363,18 +363,28 @@ func TestGetJobTypeForOwner(t *testing.T) { externalK3 := func() runtime.Object { return &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K3"}} }() + manageK4 := func() IntegrationCallbacks { + ret := dontManage + ret.IsManagingObjectsOwner = func(owner *metav1.OwnerReference) bool { return owner.Kind == "K4" } + ret.JobType = &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K4"}} + return ret + }() mgr := integrationManager{ - names: []string{"manageK1", "dontManage", "manageK2"}, + names: []string{"manageK1", "dontManage", "manageK2", "manageK4"}, integrations: map[string]IntegrationCallbacks{ "dontManage": dontManage, "manageK1": manageK1, "manageK2": manageK2, + "manageK4": manageK4, }, externalIntegrations: map[string]runtime.Object{ "externalK3": externalK3, }, } + mgr.enableIntegration("dontManage") + mgr.enableIntegration("manageK1") + mgr.enableIntegration("manageK2") cases := map[string]struct { owner *metav1.OwnerReference @@ -396,6 +406,10 @@ func TestGetJobTypeForOwner(t *testing.T) { owner: &metav1.OwnerReference{Kind: "K4"}, wantJobType: nil, }, + "K5": { + owner: &metav1.OwnerReference{Kind: "K5"}, + wantJobType: nil, + }, } for tcName, tc := range cases { diff --git a/pkg/controller/jobframework/setup.go b/pkg/controller/jobframework/setup.go index 044f3dd84b..b84440defb 100644 --- a/pkg/controller/jobframework/setup.go +++ b/pkg/controller/jobframework/setup.go @@ -83,6 +83,7 @@ func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error { if err = cb.SetupWebhook(mgr, opts...); err != nil { return fmt.Errorf("%s: unable to create webhook: %w", fwkNamePrefix, err) } + EnableIntegration(name) logger.Info("Set up controller and webhook for job framework") return nil } diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go index 891f13873f..e645004be7 100644 --- a/pkg/controller/jobs/job/job_controller_test.go +++ b/pkg/controller/jobs/job/job_controller_test.go @@ -396,6 +396,7 @@ var ( ) func TestReconciler(t *testing.T) { + defer jobframework.EnableIntegrationsForTest(FrameworkName)() baseJobWrapper := utiltestingjob.MakeJob("job", "ns"). Suspend(true). Queue("foo"). diff --git a/pkg/controller/jobs/pod/pod_controller_test.go b/pkg/controller/jobs/pod/pod_controller_test.go index 94a5485f98..0642aa88d0 100644 --- a/pkg/controller/jobs/pod/pod_controller_test.go +++ b/pkg/controller/jobs/pod/pod_controller_test.go @@ -4879,6 +4879,7 @@ func TestReconciler_ErrorFinalizingPod(t *testing.T) { } func TestIsPodOwnerManagedByQueue(t *testing.T) { + defer jobframework.EnableIntegrationsForTest("batch/job", "ray.io/raycluster")() testCases := map[string]struct { ownerReference metav1.OwnerReference wantRes bool diff --git a/pkg/controller/jobs/pod/pod_webhook_test.go b/pkg/controller/jobs/pod/pod_webhook_test.go index 2532685308..4ce5f7890a 100644 --- a/pkg/controller/jobs/pod/pod_webhook_test.go +++ b/pkg/controller/jobs/pod/pod_webhook_test.go @@ -31,6 +31,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/webhook/admission" configapi "sigs.k8s.io/kueue/apis/config/v1beta1" + "sigs.k8s.io/kueue/pkg/controller/jobframework" _ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs" _ "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob" _ "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster" @@ -64,6 +65,7 @@ func TestDefault(t *testing.T) { manageJobsWithoutQueueName bool namespaceSelector *metav1.LabelSelector podSelector *metav1.LabelSelector + enableIntegrations []string want *corev1.Pod }{ "pod with queue nil ns selector": { @@ -102,6 +104,22 @@ func TestDefault(t *testing.T) { KueueFinalizer(). Obj(), }, + "pod with owner managed by kueue (Job) while not enabled": { + initObjects: []client.Object{defaultNamespace}, + podSelector: &metav1.LabelSelector{}, + namespaceSelector: defaultNamespaceSelector, + pod: testingpod.MakePod("test-pod", defaultNamespace.Name). + Queue("test-queue"). + OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")). + Obj(), + want: testingpod.MakePod("test-pod", defaultNamespace.Name). + Queue("test-queue"). + Label("kueue.x-k8s.io/managed", "true"). + KueueSchedulingGate(). + KueueFinalizer(). + OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")). + Obj(), + }, "pod with owner managed by kueue (Job)": { initObjects: []client.Object{defaultNamespace}, podSelector: &metav1.LabelSelector{}, @@ -110,6 +128,7 @@ func TestDefault(t *testing.T) { Queue("test-queue"). OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")). Obj(), + enableIntegrations: []string{"batch/job"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")). @@ -123,6 +142,7 @@ func TestDefault(t *testing.T) { Queue("test-queue"). OwnerReference("parent-ray-cluster", rayv1.GroupVersion.WithKind("RayCluster")). Obj(), + enableIntegrations: []string{"ray.io/raycluster"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference("parent-ray-cluster", rayv1.GroupVersion.WithKind("RayCluster")). @@ -139,6 +159,7 @@ func TestDefault(t *testing.T) { schema.GroupVersionKind{Group: "kubeflow.org", Version: "v2beta1", Kind: "MPIJob"}, ). Obj(), + enableIntegrations: []string{"kubeflow.org/mpijob"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference( @@ -158,6 +179,7 @@ func TestDefault(t *testing.T) { schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "PyTorchJob"}, ). Obj(), + enableIntegrations: []string{"kubeflow.org/pytorchjob"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference( @@ -177,6 +199,7 @@ func TestDefault(t *testing.T) { schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "TFJob"}, ). Obj(), + enableIntegrations: []string{"kubeflow.org/tfjob"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference( @@ -196,6 +219,7 @@ func TestDefault(t *testing.T) { schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "XGBoostJob"}, ). Obj(), + enableIntegrations: []string{"kubeflow.org/xgboostjob"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference( @@ -215,6 +239,7 @@ func TestDefault(t *testing.T) { schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "PaddleJob"}, ). Obj(), + enableIntegrations: []string{"kubeflow.org/paddlejob"}, want: testingpod.MakePod("test-pod", defaultNamespace.Name). Queue("test-queue"). OwnerReference( @@ -261,6 +286,7 @@ func TestDefault(t *testing.T) { for name, tc := range testCases { t.Run(name, func(t *testing.T) { + defer jobframework.EnableIntegrationsForTest(tc.enableIntegrations...)() builder := utiltesting.NewClientBuilder() builder = builder.WithObjects(tc.initObjects...) cli := builder.Build() @@ -374,6 +400,7 @@ func TestGetRoleHash(t *testing.T) { } func TestValidateCreate(t *testing.T) { + defer jobframework.EnableIntegrationsForTest("batch/job")() testCases := map[string]struct { pod *corev1.Pod wantErr error @@ -476,6 +503,7 @@ func TestValidateCreate(t *testing.T) { } func TestValidateUpdate(t *testing.T) { + defer jobframework.EnableIntegrationsForTest("batch/job")() testCases := map[string]struct { oldPod *corev1.Pod newPod *corev1.Pod diff --git a/test/integration/controller/jobs/job/suite_test.go b/test/integration/controller/jobs/job/suite_test.go index 858904bb89..3a15350b0a 100644 --- a/test/integration/controller/jobs/job/suite_test.go +++ b/test/integration/controller/jobs/job/suite_test.go @@ -73,6 +73,7 @@ func managerSetup(opts ...jobframework.Option) framework.ManagerSetup { gomega.Expect(err).NotTo(gomega.HaveOccurred()) err = job.SetupWebhook(mgr, opts...) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(job.FrameworkName) } } @@ -99,6 +100,7 @@ func managerAndControllersSetup(enableScheduler bool, configuration *config.Conf gomega.Expect(err).NotTo(gomega.HaveOccurred()) err = job.SetupWebhook(mgr, opts...) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(job.FrameworkName) if enableScheduler { sched := scheduler.New(queues, cCache, mgr.GetClient(), mgr.GetEventRecorderFor(constants.AdmissionName)) diff --git a/test/integration/controller/jobs/mpijob/suite_test.go b/test/integration/controller/jobs/mpijob/suite_test.go index 577b4d8f81..dc39325242 100644 --- a/test/integration/controller/jobs/mpijob/suite_test.go +++ b/test/integration/controller/jobs/mpijob/suite_test.go @@ -73,6 +73,7 @@ func managerSetup(setupJobManager bool, opts ...jobframework.Option) framework.M gomega.Expect(err).NotTo(gomega.HaveOccurred()) err = mpijob.SetupMPIJobWebhook(mgr, opts...) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(mpijob.FrameworkName) if setupJobManager { jobReconciler := job.NewReconciler( diff --git a/test/integration/controller/jobs/pod/suite_test.go b/test/integration/controller/jobs/pod/suite_test.go index 79881f8451..bce28f8bae 100644 --- a/test/integration/controller/jobs/pod/suite_test.go +++ b/test/integration/controller/jobs/pod/suite_test.go @@ -92,6 +92,7 @@ func managerSetup(configuration *config.Configuration, opts ...jobframework.Opti opts...) err = jobReconciler.SetupWithManager(mgr) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(job.FrameworkName) cCache := cache.New(mgr.GetClient()) queues := queue.NewManager(mgr.GetClient(), cCache) diff --git a/test/integration/controller/jobs/raycluster/suite_test.go b/test/integration/controller/jobs/raycluster/suite_test.go index d862e317dc..a8e0b79e51 100644 --- a/test/integration/controller/jobs/raycluster/suite_test.go +++ b/test/integration/controller/jobs/raycluster/suite_test.go @@ -70,6 +70,7 @@ func managerSetup(opts ...jobframework.Option) framework.ManagerSetup { gomega.Expect(err).NotTo(gomega.HaveOccurred()) err = raycluster.SetupRayClusterWebhook(mgr, opts...) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(rayjob.FrameworkName) } } @@ -123,6 +124,7 @@ func managerWithRayClusterAndRayJobControllersSetup(opts ...jobframework.Option) gomega.Expect(err).NotTo(gomega.HaveOccurred()) err = rayjob.SetupRayJobWebhook(mgr, opts...) gomega.Expect(err).NotTo(gomega.HaveOccurred()) + jobframework.EnableIntegration(rayjob.FrameworkName) failedWebhook, err := webhooks.Setup(mgr) gomega.Expect(err).ToNot(gomega.HaveOccurred(), "webhook", failedWebhook) From 237531e199f16095d2cbdf8653ce00c35f25f999 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Tue, 9 Jul 2024 17:16:17 +0300 Subject: [PATCH 2/6] [jobframework] Check Enabled frameworks in unit tests --- .../jobframework/integrationmanager_test.go | 16 +++++- .../jobframework/reconciler_test.go | 1 + pkg/controller/jobframework/setup.go | 8 ++- pkg/controller/jobframework/setup_test.go | 50 +++++++++++++++++-- 4 files changed, 68 insertions(+), 7 deletions(-) diff --git a/pkg/controller/jobframework/integrationmanager_test.go b/pkg/controller/jobframework/integrationmanager_test.go index d6777a4b6c..9f8314dbe4 100644 --- a/pkg/controller/jobframework/integrationmanager_test.go +++ b/pkg/controller/jobframework/integrationmanager_test.go @@ -32,12 +32,26 @@ import ( "k8s.io/client-go/tools/record" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" + ctrlmgr "sigs.k8s.io/controller-runtime/pkg/manager" + "sigs.k8s.io/controller-runtime/pkg/reconcile" ) -func testNewReconciler(client.Client, record.EventRecorder, ...Option) JobReconcilerInterface { +type testReconciler struct{} + +func (t *testReconciler) Reconcile(context.Context, reconcile.Request) (reconcile.Result, error) { + return reconcile.Result{}, nil +} + +func (t *testReconciler) SetupWithManager(mgr ctrlmgr.Manager) error { return nil } +var _ JobReconcilerInterface = (*testReconciler)(nil) + +func testNewReconciler(client.Client, record.EventRecorder, ...Option) JobReconcilerInterface { + return &testReconciler{} +} + func testSetupWebhook(ctrl.Manager, ...Option) error { return nil } diff --git a/pkg/controller/jobframework/reconciler_test.go b/pkg/controller/jobframework/reconciler_test.go index 12fcde150b..8d2b90a012 100644 --- a/pkg/controller/jobframework/reconciler_test.go +++ b/pkg/controller/jobframework/reconciler_test.go @@ -84,6 +84,7 @@ func TestIsParentJobManaged(t *testing.T) { } for name, tc := range cases { t.Run(name, func(t *testing.T) { + defer EnableIntegrationsForTest("kubeflow.org/mpijob")() builder := utiltesting.NewClientBuilder(kubeflow.AddToScheme) if tc.parentJob != nil { builder = builder.WithObjects(tc.parentJob) diff --git a/pkg/controller/jobframework/setup.go b/pkg/controller/jobframework/setup.go index b84440defb..ebc4bf864e 100644 --- a/pkg/controller/jobframework/setup.go +++ b/pkg/controller/jobframework/setup.go @@ -45,6 +45,10 @@ var ( // until the webhooks are operating, and the webhook won't work until the // certs are all in place. func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error { + return manager.setupControllers(mgr, log, opts...) +} + +func (m *integrationManager) setupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error { options := ProcessOptions(opts...) for fwkName := range options.EnabledExternalFrameworks { @@ -52,7 +56,7 @@ func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error { return err } } - return ForEachIntegration(func(name string, cb IntegrationCallbacks) error { + return m.forEach(func(name string, cb IntegrationCallbacks) error { logger := log.WithValues("jobFrameworkName", name) fwkNamePrefix := fmt.Sprintf("jobFrameworkName %q", name) @@ -83,7 +87,7 @@ func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error { if err = cb.SetupWebhook(mgr, opts...); err != nil { return fmt.Errorf("%s: unable to create webhook: %w", fwkNamePrefix, err) } - EnableIntegration(name) + m.enableIntegration(name) logger.Info("Set up controller and webhook for job framework") return nil } diff --git a/pkg/controller/jobframework/setup_test.go b/pkg/controller/jobframework/setup_test.go index 16f8a69e0a..a472763147 100644 --- a/pkg/controller/jobframework/setup_test.go +++ b/pkg/controller/jobframework/setup_test.go @@ -43,10 +43,38 @@ import ( ) func TestSetupControllers(t *testing.T) { + aviableIntegrations := map[string]IntegrationCallbacks{ + "batch/job": IntegrationCallbacks{ + NewReconciler: testNewReconciler, + SetupWebhook: testSetupWebhook, + JobType: &batchv1.Job{}, + SetupIndexes: testSetupIndexes, + AddToScheme: testAddToScheme, + CanSupportIntegration: testCanSupportIntegration, + }, + "kubeflow.org/mpijob": IntegrationCallbacks{ + NewReconciler: testNewReconciler, + SetupWebhook: testSetupWebhook, + JobType: &kubeflow.MPIJob{}, + SetupIndexes: testSetupIndexes, + AddToScheme: testAddToScheme, + CanSupportIntegration: testCanSupportIntegration, + }, + "pod": IntegrationCallbacks{ + NewReconciler: testNewReconciler, + SetupWebhook: testSetupWebhook, + JobType: &corev1.Pod{}, + SetupIndexes: testSetupIndexes, + AddToScheme: testAddToScheme, + CanSupportIntegration: testCanSupportIntegration, + }, + } + cases := map[string]struct { - opts []Option - mapperGVKs []schema.GroupVersionKind - wantError error + opts []Option + mapperGVKs []schema.GroupVersionKind + wantError error + wantEnabledIntegrations []string }{ "setup controllers succeed": { opts: []Option{ @@ -60,6 +88,7 @@ func TestSetupControllers(t *testing.T) { batchv1.SchemeGroupVersion.WithKind("Job"), kubeflow.SchemeGroupVersionKind, }, + wantEnabledIntegrations: []string{"batch/job", "kubeflow.org/mpijob"}, }, "mapper doesn't have kubeflow.org/mpijob, but no error occur": { opts: []Option{ @@ -68,10 +97,19 @@ func TestSetupControllers(t *testing.T) { mapperGVKs: []schema.GroupVersionKind{ batchv1.SchemeGroupVersion.WithKind("Job"), }, + wantEnabledIntegrations: []string{"batch/job"}, }, } for name, tc := range cases { t.Run(name, func(t *testing.T) { + manager := integrationManager{} + for name, cbs := range aviableIntegrations { + err := manager.register(name, cbs) + if err != nil { + t.Fatalf("Unexpected error while registering %q: %s", name, err) + } + } + _, logger := utiltesting.ContextWithLog(t) k8sClient := utiltesting.NewClientBuilder(jobset.AddToScheme, kubeflow.AddToScheme, kftraining.AddToScheme, rayv1.AddToScheme).Build() @@ -97,10 +135,14 @@ func TestSetupControllers(t *testing.T) { t.Fatalf("Failed to setup manager: %v", err) } - gotError := SetupControllers(mgr, logger, tc.opts...) + gotError := manager.setupControllers(mgr, logger, tc.opts...) if diff := cmp.Diff(tc.wantError, gotError, cmpopts.EquateErrors()); len(diff) != 0 { t.Errorf("Unexpected error from SetupControllers (-want,+got):\n%s", diff) } + + if diff := cmp.Diff(tc.wantEnabledIntegrations, manager.enabledIntegrations.SortedList()); len(diff) != 0 { + t.Errorf("Unexpected enabled integrations (-want,+got):\n%s", diff) + } }) } } From c8ba8b520d12c524cd59a17074ddb46b0614ddb6 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Wed, 10 Jul 2024 16:37:08 +0300 Subject: [PATCH 3/6] Review Remarks --- pkg/controller/jobframework/setup_test.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pkg/controller/jobframework/setup_test.go b/pkg/controller/jobframework/setup_test.go index a472763147..03f59b0db2 100644 --- a/pkg/controller/jobframework/setup_test.go +++ b/pkg/controller/jobframework/setup_test.go @@ -43,8 +43,8 @@ import ( ) func TestSetupControllers(t *testing.T) { - aviableIntegrations := map[string]IntegrationCallbacks{ - "batch/job": IntegrationCallbacks{ + availableIntegrations := map[string]IntegrationCallbacks{ + "batch/job": { NewReconciler: testNewReconciler, SetupWebhook: testSetupWebhook, JobType: &batchv1.Job{}, @@ -52,7 +52,7 @@ func TestSetupControllers(t *testing.T) { AddToScheme: testAddToScheme, CanSupportIntegration: testCanSupportIntegration, }, - "kubeflow.org/mpijob": IntegrationCallbacks{ + "kubeflow.org/mpijob": { NewReconciler: testNewReconciler, SetupWebhook: testSetupWebhook, JobType: &kubeflow.MPIJob{}, @@ -60,7 +60,7 @@ func TestSetupControllers(t *testing.T) { AddToScheme: testAddToScheme, CanSupportIntegration: testCanSupportIntegration, }, - "pod": IntegrationCallbacks{ + "pod": { NewReconciler: testNewReconciler, SetupWebhook: testSetupWebhook, JobType: &corev1.Pod{}, @@ -103,7 +103,7 @@ func TestSetupControllers(t *testing.T) { for name, tc := range cases { t.Run(name, func(t *testing.T) { manager := integrationManager{} - for name, cbs := range aviableIntegrations { + for name, cbs := range availableIntegrations { err := manager.register(name, cbs) if err != nil { t.Fatalf("Unexpected error while registering %q: %s", name, err) From b6105150e0a544ba25d85602667a79636dd83f87 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Thu, 11 Jul 2024 07:18:06 +0300 Subject: [PATCH 4/6] Review Remarks --- pkg/controller/jobframework/integrationmanager.go | 4 +++- pkg/controller/jobframework/reconciler_test.go | 2 +- pkg/controller/jobs/job/job_controller_test.go | 2 +- pkg/controller/jobs/pod/pod_controller_test.go | 2 +- pkg/controller/jobs/pod/pod_webhook_test.go | 6 +++--- 5 files changed, 9 insertions(+), 7 deletions(-) diff --git a/pkg/controller/jobframework/integrationmanager.go b/pkg/controller/jobframework/integrationmanager.go index 2b0cf61da4..dd634057bf 100644 --- a/pkg/controller/jobframework/integrationmanager.go +++ b/pkg/controller/jobframework/integrationmanager.go @@ -21,6 +21,7 @@ import ( "errors" "fmt" "sort" + "testing" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" @@ -204,7 +205,8 @@ func EnableIntegration(name string) { // EnableIntegrationsForTest - should be used only in tests // Mark the frameworks identified by names and return a revert function. -func EnableIntegrationsForTest(names ...string) func() { +func EnableIntegrationsForTest(tb testing.TB, names ...string) func() { + tb.Helper() old := manager.enabledIntegrations.Clone() for _, name := range names { manager.enableIntegration(name) diff --git a/pkg/controller/jobframework/reconciler_test.go b/pkg/controller/jobframework/reconciler_test.go index 8d2b90a012..11dc29b399 100644 --- a/pkg/controller/jobframework/reconciler_test.go +++ b/pkg/controller/jobframework/reconciler_test.go @@ -84,7 +84,7 @@ func TestIsParentJobManaged(t *testing.T) { } for name, tc := range cases { t.Run(name, func(t *testing.T) { - defer EnableIntegrationsForTest("kubeflow.org/mpijob")() + defer EnableIntegrationsForTest(t, "kubeflow.org/mpijob")() builder := utiltesting.NewClientBuilder(kubeflow.AddToScheme) if tc.parentJob != nil { builder = builder.WithObjects(tc.parentJob) diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go index e645004be7..29aa874607 100644 --- a/pkg/controller/jobs/job/job_controller_test.go +++ b/pkg/controller/jobs/job/job_controller_test.go @@ -396,7 +396,7 @@ var ( ) func TestReconciler(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(FrameworkName)() + defer jobframework.EnableIntegrationsForTest(t, FrameworkName)() baseJobWrapper := utiltestingjob.MakeJob("job", "ns"). Suspend(true). Queue("foo"). diff --git a/pkg/controller/jobs/pod/pod_controller_test.go b/pkg/controller/jobs/pod/pod_controller_test.go index 0642aa88d0..212937b3be 100644 --- a/pkg/controller/jobs/pod/pod_controller_test.go +++ b/pkg/controller/jobs/pod/pod_controller_test.go @@ -4879,7 +4879,7 @@ func TestReconciler_ErrorFinalizingPod(t *testing.T) { } func TestIsPodOwnerManagedByQueue(t *testing.T) { - defer jobframework.EnableIntegrationsForTest("batch/job", "ray.io/raycluster")() + defer jobframework.EnableIntegrationsForTest(t, "batch/job", "ray.io/raycluster")() testCases := map[string]struct { ownerReference metav1.OwnerReference wantRes bool diff --git a/pkg/controller/jobs/pod/pod_webhook_test.go b/pkg/controller/jobs/pod/pod_webhook_test.go index 4ce5f7890a..3b0cd5d159 100644 --- a/pkg/controller/jobs/pod/pod_webhook_test.go +++ b/pkg/controller/jobs/pod/pod_webhook_test.go @@ -286,7 +286,7 @@ func TestDefault(t *testing.T) { for name, tc := range testCases { t.Run(name, func(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(tc.enableIntegrations...)() + defer jobframework.EnableIntegrationsForTest(t, tc.enableIntegrations...)() builder := utiltesting.NewClientBuilder() builder = builder.WithObjects(tc.initObjects...) cli := builder.Build() @@ -400,7 +400,7 @@ func TestGetRoleHash(t *testing.T) { } func TestValidateCreate(t *testing.T) { - defer jobframework.EnableIntegrationsForTest("batch/job")() + defer jobframework.EnableIntegrationsForTest(t, "batch/job")() testCases := map[string]struct { pod *corev1.Pod wantErr error @@ -503,7 +503,7 @@ func TestValidateCreate(t *testing.T) { } func TestValidateUpdate(t *testing.T) { - defer jobframework.EnableIntegrationsForTest("batch/job")() + defer jobframework.EnableIntegrationsForTest(t, "batch/job")() testCases := map[string]struct { oldPod *corev1.Pod newPod *corev1.Pod From f5866f60efa55acb32dc59cbb1ed109482e64c4b Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Fri, 12 Jul 2024 09:33:37 +0300 Subject: [PATCH 5/6] Review Remarks --- pkg/controller/jobframework/integrationmanager_test.go | 6 +++--- pkg/controller/jobs/pod/pod_controller_test.go | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pkg/controller/jobframework/integrationmanager_test.go b/pkg/controller/jobframework/integrationmanager_test.go index 9f8314dbe4..9375773e04 100644 --- a/pkg/controller/jobframework/integrationmanager_test.go +++ b/pkg/controller/jobframework/integrationmanager_test.go @@ -377,7 +377,7 @@ func TestGetJobTypeForOwner(t *testing.T) { externalK3 := func() runtime.Object { return &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K3"}} }() - manageK4 := func() IntegrationCallbacks { + disabledK4 := func() IntegrationCallbacks { ret := dontManage ret.IsManagingObjectsOwner = func(owner *metav1.OwnerReference) bool { return owner.Kind == "K4" } ret.JobType = &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K4"}} @@ -385,12 +385,12 @@ func TestGetJobTypeForOwner(t *testing.T) { }() mgr := integrationManager{ - names: []string{"manageK1", "dontManage", "manageK2", "manageK4"}, + names: []string{"manageK1", "dontManage", "manageK2", "disabledK4"}, integrations: map[string]IntegrationCallbacks{ "dontManage": dontManage, "manageK1": manageK1, "manageK2": manageK2, - "manageK4": manageK4, + "disabledK4": disabledK4, }, externalIntegrations: map[string]runtime.Object{ "externalK3": externalK3, diff --git a/pkg/controller/jobs/pod/pod_controller_test.go b/pkg/controller/jobs/pod/pod_controller_test.go index 212937b3be..127589dc3a 100644 --- a/pkg/controller/jobs/pod/pod_controller_test.go +++ b/pkg/controller/jobs/pod/pod_controller_test.go @@ -4879,7 +4879,7 @@ func TestReconciler_ErrorFinalizingPod(t *testing.T) { } func TestIsPodOwnerManagedByQueue(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(t, "batch/job", "ray.io/raycluster")() + t.Cleanup(jobframework.EnableIntegrationsForTest(t, "batch/job", "ray.io/raycluster")) testCases := map[string]struct { ownerReference metav1.OwnerReference wantRes bool From a31adaeaf2ec2bff2908c48c9e0e39b46e80a789 Mon Sep 17 00:00:00 2001 From: Traian Schiau Date: Fri, 12 Jul 2024 17:44:36 +0300 Subject: [PATCH 6/6] Review Remarks --- pkg/controller/jobframework/reconciler_test.go | 2 +- pkg/controller/jobs/job/job_controller_test.go | 2 +- pkg/controller/jobs/pod/pod_webhook_test.go | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pkg/controller/jobframework/reconciler_test.go b/pkg/controller/jobframework/reconciler_test.go index 11dc29b399..e2d1b6ffb1 100644 --- a/pkg/controller/jobframework/reconciler_test.go +++ b/pkg/controller/jobframework/reconciler_test.go @@ -84,7 +84,7 @@ func TestIsParentJobManaged(t *testing.T) { } for name, tc := range cases { t.Run(name, func(t *testing.T) { - defer EnableIntegrationsForTest(t, "kubeflow.org/mpijob")() + t.Cleanup(EnableIntegrationsForTest(t, "kubeflow.org/mpijob")) builder := utiltesting.NewClientBuilder(kubeflow.AddToScheme) if tc.parentJob != nil { builder = builder.WithObjects(tc.parentJob) diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go index 29aa874607..15bae696be 100644 --- a/pkg/controller/jobs/job/job_controller_test.go +++ b/pkg/controller/jobs/job/job_controller_test.go @@ -396,7 +396,7 @@ var ( ) func TestReconciler(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(t, FrameworkName)() + t.Cleanup(jobframework.EnableIntegrationsForTest(t, FrameworkName)) baseJobWrapper := utiltestingjob.MakeJob("job", "ns"). Suspend(true). Queue("foo"). diff --git a/pkg/controller/jobs/pod/pod_webhook_test.go b/pkg/controller/jobs/pod/pod_webhook_test.go index 3b0cd5d159..66f2c0bb21 100644 --- a/pkg/controller/jobs/pod/pod_webhook_test.go +++ b/pkg/controller/jobs/pod/pod_webhook_test.go @@ -286,7 +286,7 @@ func TestDefault(t *testing.T) { for name, tc := range testCases { t.Run(name, func(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(t, tc.enableIntegrations...)() + t.Cleanup(jobframework.EnableIntegrationsForTest(t, tc.enableIntegrations...)) builder := utiltesting.NewClientBuilder() builder = builder.WithObjects(tc.initObjects...) cli := builder.Build() @@ -400,7 +400,7 @@ func TestGetRoleHash(t *testing.T) { } func TestValidateCreate(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(t, "batch/job")() + t.Cleanup(jobframework.EnableIntegrationsForTest(t, "batch/job")) testCases := map[string]struct { pod *corev1.Pod wantErr error @@ -503,7 +503,7 @@ func TestValidateCreate(t *testing.T) { } func TestValidateUpdate(t *testing.T) { - defer jobframework.EnableIntegrationsForTest(t, "batch/job")() + t.Cleanup(jobframework.EnableIntegrationsForTest(t, "batch/job")) testCases := map[string]struct { oldPod *corev1.Pod newPod *corev1.Pod