diff --git a/pkg/cache/application_test.go b/pkg/cache/application_test.go index 8f52f33de..418b21940 100644 --- a/pkg/cache/application_test.go +++ b/pkg/cache/application_test.go @@ -139,6 +139,7 @@ func TestFailApplication(t *testing.T) { rt.time++ } events.SetRecorder(mr) + defer events.SetRecorder(events.NewMockedRecorder()) resources := make(map[v1.ResourceName]resource.Quantity) containers := make([]v1.Container, 0) containers = append(containers, v1.Container{ @@ -200,8 +201,6 @@ func TestFailApplication(t *testing.T) { assert.NilError(t, err) assertAppState(t, app2, ApplicationStates().Failed, 3*time.Second) assert.Equal(t, rt.time, int64(0)) - // Test over, set Recorder back fake type - events.SetRecorder(k8sEvents.NewFakeRecorder(1024)) } func TestSetUnallocatedPodsToFailedWhenFailApplication(t *testing.T) { @@ -228,6 +227,7 @@ func TestSetUnallocatedPodsToFailedWhenFailApplication(t *testing.T) { // set Recorder to mocked type mr := events.NewMockedRecorder() events.SetRecorder(mr) + defer events.SetRecorder(events.NewMockedRecorder()) resources := make(map[v1.ResourceName]resource.Quantity) containers := make([]v1.Container, 0) containers = append(containers, v1.Container{ @@ -306,8 +306,6 @@ func TestSetUnallocatedPodsToFailedWhenFailApplication(t *testing.T) { assert.NilError(t, err) assert.Equal(t, newPod3.Status.Phase, v1.PodFailed, 3*time.Second) assert.Equal(t, newPod3.Status.Reason, constants.ApplicationInsufficientResourcesFailure, 3*time.Second) - // Test over, set Recorder back fake type - events.SetRecorder(k8sEvents.NewFakeRecorder(1024)) } func TestSetUnallocatedPodsToFailedWhenRejectApplication(t *testing.T) { @@ -334,6 +332,7 @@ func TestSetUnallocatedPodsToFailedWhenRejectApplication(t *testing.T) { // set Recorder to mocked type mr := events.NewMockedRecorder() events.SetRecorder(mr) + defer events.SetRecorder(events.NewMockedRecorder()) resources := make(map[v1.ResourceName]resource.Quantity) containers := make([]v1.Container, 0) containers = append(containers, v1.Container{ @@ -404,8 +403,6 @@ func TestSetUnallocatedPodsToFailedWhenRejectApplication(t *testing.T) { assert.NilError(t, err) assert.Equal(t, newPod2.Status.Phase, v1.PodFailed, 3*time.Second) assert.Equal(t, newPod2.Status.Reason, constants.ApplicationRejectedFailure, 3*time.Second) - // Test over, set Recorder back fake type - events.SetRecorder(k8sEvents.NewFakeRecorder(1024)) } func TestReleaseAppAllocation(t *testing.T) { @@ -1110,10 +1107,9 @@ func TestGetPlaceholderTasks(t *testing.T) { func TestPlaceholderTimeoutEvents(t *testing.T) { context := initContextForTest() - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) pod1 := v1.Pod{ TypeMeta: apis.TypeMeta{ @@ -1213,10 +1209,9 @@ func TestApplication_onReservationStateChange(t *testing.T) { dispatcher.Start() defer dispatcher.Stop() - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) app := NewApplication(appID, "root.a", "testuser", testGroups, map[string]string{}, newMockSchedulerAPI()) context.addApplicationToContext(app) diff --git a/pkg/cache/context_test.go b/pkg/cache/context_test.go index 91f48e60a..a3a6fc92b 100644 --- a/pkg/cache/context_test.go +++ b/pkg/cache/context_test.go @@ -1059,10 +1059,9 @@ func TestGetTask(t *testing.T) { } func TestNodeEventFailsPublishingWithoutNode(t *testing.T) { - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) context := initContextForTest() eventRecords := make([]*si.EventRecord, 0) @@ -1087,10 +1086,9 @@ func TestNodeEventFailsPublishingWithoutNode(t *testing.T) { } func TestNodeEventPublishedCorrectly(t *testing.T) { - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) context, apiProvider := initContextAndAPIProviderForTest() dispatcher.Start() @@ -1147,10 +1145,9 @@ func TestNodeEventPublishedCorrectly(t *testing.T) { } func TestFilteredEventsNotPublished(t *testing.T) { - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) context, apiProvider := initContextAndAPIProviderForTest() dispatcher.Start() @@ -1231,10 +1228,10 @@ func TestFilteredEventsNotPublished(t *testing.T) { } func TestPublishEventsWithNotExistingAsk(t *testing.T) { - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) + context := initContextForTest() context.AddApplication(&AddApplicationRequest{ Metadata: ApplicationMetadata{ @@ -1271,10 +1268,10 @@ func TestPublishEventsWithNotExistingAsk(t *testing.T) { } func TestPublishEventsCorrectly(t *testing.T) { - recorder, ok := events.GetRecorder().(*k8sEvents.FakeRecorder) - if !ok { - t.Fatal("the EventRecorder is expected to be of type FakeRecorder") - } + recorder := k8sEvents.NewFakeRecorder(1024) + events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) + context := initContextForTest() // create fake application and task diff --git a/pkg/cache/scheduler_callback_test.go b/pkg/cache/scheduler_callback_test.go index d98f5f6f7..8b0db9d98 100644 --- a/pkg/cache/scheduler_callback_test.go +++ b/pkg/cache/scheduler_callback_test.go @@ -248,6 +248,7 @@ func TestUpdateApplication_Rejected(t *testing.T) { NewPlaceholderManager(context.apiProvider.GetAPIs()) recorder := k8sEvents.NewFakeRecorder(1024) events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) err := callback.UpdateApplication(&si.ApplicationResponse{ Rejected: []*si.RejectedApplication{ @@ -361,6 +362,7 @@ func testUpdateApplicationFailure(t *testing.T, state string) { NewPlaceholderManager(context.apiProvider.GetAPIs()) recorder := k8sEvents.NewFakeRecorder(1024) events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) err := callback.UpdateApplication(&si.ApplicationResponse{ Updated: []*si.UpdatedApplication{ @@ -459,6 +461,7 @@ func TestSendEvent(t *testing.T) { callback, _ := initCallbackTest(t, false, false) recorder := k8sEvents.NewFakeRecorder(1024) events.SetRecorder(recorder) + defer events.SetRecorder(events.NewMockedRecorder()) defer dispatcher.UnregisterAllEventHandlers() defer dispatcher.Stop() diff --git a/pkg/cache/task_test.go b/pkg/cache/task_test.go index 864cd7e87..415afa17d 100644 --- a/pkg/cache/task_test.go +++ b/pkg/cache/task_test.go @@ -23,19 +23,16 @@ import ( "time" "gotest.tools/v3/assert" - v1 "k8s.io/api/core/v1" schedulingv1 "k8s.io/api/scheduling/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - k8sEvents "k8s.io/client-go/tools/events" "github.com/apache/yunikorn-k8shim/pkg/client" "github.com/apache/yunikorn-k8shim/pkg/common/constants" "github.com/apache/yunikorn-k8shim/pkg/common/events" "github.com/apache/yunikorn-k8shim/pkg/common/utils" "github.com/apache/yunikorn-k8shim/pkg/locking" - "github.com/apache/yunikorn-scheduler-interface/lib/go/si" ) @@ -527,6 +524,7 @@ func TestHandleSubmitTaskEvent(t *testing.T) { rt.time++ } events.SetRecorder(mr) + defer events.SetRecorder(events.NewMockedRecorder()) resources := make(map[v1.ResourceName]resource.Quantity) containers := make([]v1.Container, 0) containers = append(containers, v1.Container{ @@ -604,9 +602,6 @@ func TestHandleSubmitTaskEvent(t *testing.T) { assert.Assert(t, allocRequest.Allocations[0].PreemptionPolicy != nil) assert.Assert(t, !allocRequest.Allocations[0].PreemptionPolicy.AllowPreemptSelf) assert.Assert(t, allocRequest.Allocations[0].PreemptionPolicy.AllowPreemptOther) - - // Test over, set Recorder back fake type - events.SetRecorder(k8sEvents.NewFakeRecorder(1024)) } func TestSimultaneousTaskCompleteAndAllocate(t *testing.T) {