[YUNIKORN-2971] Use FakeRecorder and MockedRecorder properly in the t…

…ests
apache · Nov 12, 2024 · 1b81428 · 1b81428
1 parent 8fa300f
commit 1b81428
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 40 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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) {