This is an automated email from the ASF dual-hosted git repository.
mani pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/yunikorn-core.git
The following commit(s) were added to refs/heads/master by this push:
new d23cfe5e [YUNIKORN-2375] Refactor mock callback plugin (#782)
d23cfe5e is described below
commit d23cfe5ecd697060f739288c10ade0ed5d4c1fd7
Author: Wilfred Spiegelenburg <[email protected]>
AuthorDate: Fri Feb 2 12:32:51 2024 +0530
[YUNIKORN-2375] Refactor mock callback plugin (#782)
Move the mock code for the resource manager callback out of the test
directory into its own top level directory.
Removes imports of scheduler test code outside the scheduler package.
Closes: #782
Signed-off-by: Manikandan R <[email protected]>
---
pkg/events/event_publisher_test.go | 55 ++------
pkg/examples/simple_example.go | 6 +-
.../container_state_updater.go} | 33 +++--
.../plugins_test.go => mock/event_plugin.go} | 45 +++++--
.../predicate_plugin.go} | 15 +--
pkg/mock/preemption_predicate_plugin.go | 144 +++++++++++++++++++++
.../mock_rm_callback.go => mock/rm_callback.go} | 21 +--
pkg/plugins/plugins_test.go | 8 +-
pkg/scheduler/objects/preemption_test.go | 128 +-----------------
pkg/scheduler/partition_test.go | 5 +-
pkg/scheduler/tests/mock_rm_callback_test.go | 5 +-
pkg/scheduler/tests/plugin_test.go | 26 +---
12 files changed, 245 insertions(+), 246 deletions(-)
diff --git a/pkg/events/event_publisher_test.go
b/pkg/events/event_publisher_test.go
index 88d2b259..d83ffae8 100644
--- a/pkg/events/event_publisher_test.go
+++ b/pkg/events/event_publisher_test.go
@@ -19,59 +19,17 @@
package events
import (
- "fmt"
- "sync"
"testing"
"time"
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-core/pkg/plugins"
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type mockEventPlugin struct {
- tests.MockResourceManagerCallback
- records chan *si.EventRecord
-
- sync.Mutex
-}
-
-// create and register mocked event plugin
-func createEventPluginForTest() (*mockEventPlugin, error) {
- eventPlugin := mockEventPlugin{
- records: make(chan *si.EventRecord, 3),
- }
- plugins.RegisterSchedulerPlugin(&eventPlugin)
- if plugins.GetResourceManagerCallbackPlugin() == nil {
- return nil, fmt.Errorf("event plugin is not registered")
- }
- return &eventPlugin, nil
-}
-
-func (ep *mockEventPlugin) SendEvent(events []*si.EventRecord) {
- ep.Lock()
- defer ep.Unlock()
-
- for _, event := range events {
- ep.records <- event
- }
-}
-
-func (ep *mockEventPlugin) getNextEventRecord() *si.EventRecord {
- ep.Lock()
- defer ep.Unlock()
-
- select {
- case record := <-ep.records:
- return record
- default:
- return nil
- }
-}
-
// creating a Publisher with nil store should still provide a non-nil object
func TestCreateShimPublisher(t *testing.T) {
publisher := CreateShimPublisher(nil)
@@ -112,8 +70,11 @@ func TestNoFillWithoutEventPluginRegistered(t *testing.T) {
// we push an event to the publisher, and check that the same event
// is published by observing the mocked EventPlugin
func TestPublisherSendsEvent(t *testing.T) {
- eventPlugin, err := createEventPluginForTest()
- assert.NilError(t, err, "could not create event plugin for test")
+ eventPlugin := mock.NewEventPlugin()
+ plugins.RegisterSchedulerPlugin(eventPlugin)
+ if plugins.GetResourceManagerCallbackPlugin() == nil {
+ t.Fatal("could not register event plugin for test")
+ }
store := newEventStore()
publisher := CreateShimPublisher(store)
@@ -131,8 +92,8 @@ func TestPublisherSendsEvent(t *testing.T) {
store.Store(event)
var eventFromPlugin *si.EventRecord
- err = common.WaitForCondition(func() bool {
- eventFromPlugin = eventPlugin.getNextEventRecord()
+ err := common.WaitForCondition(func() bool {
+ eventFromPlugin = eventPlugin.GetNextEventRecord()
return eventFromPlugin != nil
}, time.Millisecond, time.Second)
assert.NilError(t, err, "event was not received in time: %v", err)
diff --git a/pkg/examples/simple_example.go b/pkg/examples/simple_example.go
index dd8db616..077be0d6 100644
--- a/pkg/examples/simple_example.go
+++ b/pkg/examples/simple_example.go
@@ -22,12 +22,12 @@ import (
"sync"
"github.com/apache/yunikorn-core/pkg/entrypoint"
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
type exampleRMCallback struct {
- tests.MockResourceManagerCallback
+ mock.ResourceManagerCallback
acceptedApplications map[string]bool
rejectedApplications map[string]bool
acceptedNodes map[string]bool
@@ -205,7 +205,7 @@ partitions:
panic(err)
}
- // Refer to mock_rm_callback.go:109
+ // Refer to mock/rm_callback.go:109
// You need to check app accepted by scheduler before proceed.
// Send request
diff --git a/pkg/plugins/plugins_test.go b/pkg/mock/container_state_updater.go
similarity index 52%
copy from pkg/plugins/plugins_test.go
copy to pkg/mock/container_state_updater.go
index 9c5fef12..e5890032 100644
--- a/pkg/plugins/plugins_test.go
+++ b/pkg/mock/container_state_updater.go
@@ -16,22 +16,33 @@
limitations under the License.
*/
-package plugins
+package mock
import (
- "testing"
+ "sync"
- "gotest.tools/v3/assert"
-
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
+ "github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type FakeResourceManagerCallback struct {
- tests.MockResourceManagerCallback
+type ContainerStateUpdater struct {
+ ResourceManagerCallback
+ sentUpdate *si.UpdateContainerSchedulingStateRequest
+ sync.RWMutex
+}
+
+func (m *ContainerStateUpdater) UpdateContainerSchedulingState(request
*si.UpdateContainerSchedulingStateRequest) {
+ m.Lock()
+ defer m.Unlock()
+ m.sentUpdate = request
+}
+
+func (m *ContainerStateUpdater) GetContainerUpdateRequest()
*si.UpdateContainerSchedulingStateRequest {
+ m.RLock()
+ defer m.RUnlock()
+ return m.sentUpdate
}
-func TestRegisterPlugins(t *testing.T) {
- plugins = SchedulerPlugins{}
- RegisterSchedulerPlugin(&FakeResourceManagerCallback{})
- assert.Assert(t, GetResourceManagerCallbackPlugin() != nil,
"ResourceManagerCallbackPlugin plugin should have been registered")
+// NewContainerStateUpdater returns a mock that can allows retrieval of the
update that was sent.
+func NewContainerStateUpdater() *ContainerStateUpdater {
+ return &ContainerStateUpdater{}
}
diff --git a/pkg/plugins/plugins_test.go b/pkg/mock/event_plugin.go
similarity index 55%
copy from pkg/plugins/plugins_test.go
copy to pkg/mock/event_plugin.go
index 9c5fef12..d4584add 100644
--- a/pkg/plugins/plugins_test.go
+++ b/pkg/mock/event_plugin.go
@@ -16,22 +16,45 @@
limitations under the License.
*/
-package plugins
+package mock
import (
- "testing"
+ "sync"
- "gotest.tools/v3/assert"
-
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
+ "github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type FakeResourceManagerCallback struct {
- tests.MockResourceManagerCallback
+type EventPlugin struct {
+ ResourceManagerCallback
+ records chan *si.EventRecord
+
+ sync.Mutex
+}
+
+func (m *EventPlugin) SendEvent(events []*si.EventRecord) {
+ m.Lock()
+ defer m.Unlock()
+
+ for _, event := range events {
+ m.records <- event
+ }
+}
+
+func (m *EventPlugin) GetNextEventRecord() *si.EventRecord {
+ m.Lock()
+ defer m.Unlock()
+
+ select {
+ case record := <-m.records:
+ return record
+ default:
+ return nil
+ }
}
-func TestRegisterPlugins(t *testing.T) {
- plugins = SchedulerPlugins{}
- RegisterSchedulerPlugin(&FakeResourceManagerCallback{})
- assert.Assert(t, GetResourceManagerCallbackPlugin() != nil,
"ResourceManagerCallbackPlugin plugin should have been registered")
+// NewEventPlugin creates a mocked event plugin
+func NewEventPlugin() *EventPlugin {
+ return &EventPlugin{
+ records: make(chan *si.EventRecord, 3),
+ }
}
diff --git a/pkg/scheduler/mock_plugin_test.go b/pkg/mock/predicate_plugin.go
similarity index 82%
rename from pkg/scheduler/mock_plugin_test.go
rename to pkg/mock/predicate_plugin.go
index c5f0f87e..27aa3822 100644
--- a/pkg/scheduler/mock_plugin_test.go
+++ b/pkg/mock/predicate_plugin.go
@@ -16,7 +16,7 @@
limitations under the License.
*/
-package scheduler
+package mock
import (
"fmt"
@@ -24,17 +24,16 @@ import (
"go.uber.org/zap"
"github.com/apache/yunikorn-core/pkg/log"
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type fakePredicatePlugin struct {
- tests.MockResourceManagerCallback
+type PredicatePlugin struct {
+ ResourceManagerCallback
mustFail bool
nodes map[string]int
}
-func (f *fakePredicatePlugin) Predicates(args *si.PredicatesArgs) error {
+func (f *PredicatePlugin) Predicates(args *si.PredicatesArgs) error {
if f.mustFail {
log.Log(log.Test).Info("fake predicate plugin fail: must fail
set")
return fmt.Errorf("fake predicate plugin failed")
@@ -58,12 +57,12 @@ func (f *fakePredicatePlugin) Predicates(args
*si.PredicatesArgs) error {
return nil
}
-// A fake predicate plugin that can either always fail or fail based on the
node that is checked.
+// NewPredicatePlugin returns a mock that can either always fail or fail based
on the node that is checked.
// mustFail will cause the predicate check to always fail
// nodes allows specifying which node to fail for which check using the nodeID:
// possible values: -1 fail reserve, 0 fail always, 1 fail alloc (defaults to
always)
-func newFakePredicatePlugin(mustFail bool, nodes map[string]int)
*fakePredicatePlugin {
- return &fakePredicatePlugin{
+func NewPredicatePlugin(mustFail bool, nodes map[string]int) *PredicatePlugin {
+ return &PredicatePlugin{
mustFail: mustFail,
nodes: nodes,
}
diff --git a/pkg/mock/preemption_predicate_plugin.go
b/pkg/mock/preemption_predicate_plugin.go
new file mode 100644
index 00000000..af47c3ac
--- /dev/null
+++ b/pkg/mock/preemption_predicate_plugin.go
@@ -0,0 +1,144 @@
+/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+
+package mock
+
+import (
+ "errors"
+ "fmt"
+
+ "github.com/apache/yunikorn-scheduler-interface/lib/go/si"
+)
+
+type PreemptionPredicatePlugin struct {
+ ResourceManagerCallback
+ reservations map[string]string
+ allocations map[string]string
+ preemptions []Preemption
+ errHolder *errHolder
+}
+
+type Preemption struct {
+ expectedAllocationKey string
+ expectedNodeID string
+ expectedAllocationKeys []string
+ expectedStartIndex int32
+ success bool
+ index int32
+}
+
+type errHolder struct {
+ err error
+}
+
+func (m *PreemptionPredicatePlugin) Predicates(args *si.PredicatesArgs) error {
+ if args.Allocate {
+ nodeID, ok := m.allocations[args.AllocationKey]
+ if !ok {
+ return errors.New("no allocation found")
+ }
+ if nodeID != args.NodeID {
+ return errors.New("wrong node")
+ }
+ return nil
+ } else {
+ nodeID, ok := m.reservations[args.AllocationKey]
+ if !ok {
+ return errors.New("no allocation found")
+ }
+ if nodeID != args.NodeID {
+ return errors.New("wrong node")
+ }
+ return nil
+ }
+}
+
+func (m *PreemptionPredicatePlugin) PreemptionPredicates(args
*si.PreemptionPredicatesArgs) *si.PreemptionPredicatesResponse {
+ result := &si.PreemptionPredicatesResponse{
+ Success: false,
+ Index: -1,
+ }
+ for _, preemption := range m.preemptions {
+ if preemption.expectedAllocationKey != args.AllocationKey {
+ continue
+ }
+ if preemption.expectedNodeID != args.NodeID {
+ continue
+ }
+ if preemption.expectedStartIndex != args.StartIndex {
+ m.errHolder.err = fmt.Errorf("unexpected start index
exepected=%d, actual=%d, allocationKey=%s",
+ preemption.expectedStartIndex, args.StartIndex,
args.AllocationKey)
+ return result
+ }
+ if len(preemption.expectedAllocationKeys) !=
len(args.PreemptAllocationKeys) {
+ m.errHolder.err = fmt.Errorf("unexpected alloc key
length expected=%d, actual=%d, allocationKey=%s",
+ len(preemption.expectedAllocationKeys),
len(args.PreemptAllocationKeys), args.AllocationKey)
+ return result
+ }
+ for idx, key := range preemption.expectedAllocationKeys {
+ if args.PreemptAllocationKeys[idx] != key {
+ m.errHolder.err = fmt.Errorf("unexpected
preempt alloc key expected=%s, actual=%s, index=%d, allocationKey=%s",
+ args.PreemptAllocationKeys[idx], key,
idx, args.AllocationKey)
+ return result
+ }
+ }
+ m.errHolder = &errHolder{}
+ result.Success = preemption.success
+ result.Index = preemption.index
+ return result
+ }
+ m.errHolder.err = fmt.Errorf("mo match found allocationKey=%s,
nodeID=%s", args.AllocationKey, args.NodeID)
+ return result
+}
+
+// GetPredicateError returns the error set by the preemption predicate check
that failed.
+// Returns a nil error on success.
+func (m *PreemptionPredicatePlugin) GetPredicateError() error {
+ return m.errHolder.err
+}
+
+// NewPreemptionPredicatePlugin returns a mock plugin that can handle multiple
predicate scenarios.
+// reservations: provide a list of allocations and node IDs for which the
reservation predicate succeeds
+// allocs: provide a list of allocations and node IDs for which the allocation
predicate succeeds
+// preempt: a slice of preemption scenarios configured for the plugin to check
+func NewPreemptionPredicatePlugin(reservations, allocs map[string]string,
preempt []Preemption) *PreemptionPredicatePlugin {
+ return &PreemptionPredicatePlugin{
+ reservations: reservations,
+ allocations: allocs,
+ preemptions: preempt,
+ errHolder: &errHolder{},
+ }
+}
+
+// NewPreemption returns a preemption scenario
+// success: overall success state
+// allocKey: allocation key for this preemption scenario
+// nodeID: node for this preemption scenario
+// expectedPreemptions: the allocations that should be in the preemption list
+// start: the point at which to start the checks
+// index: the index into the expectedPreemptions to return on success
+func NewPreemption(success bool, allocKey, nodeID string, expectedPreemptions
[]string, start, index int32) Preemption {
+ return Preemption{
+ expectedAllocationKey: allocKey,
+ expectedNodeID: nodeID,
+ expectedAllocationKeys: expectedPreemptions,
+ expectedStartIndex: start,
+ success: success,
+ index: index,
+ }
+}
diff --git a/pkg/scheduler/tests/mock_rm_callback.go b/pkg/mock/rm_callback.go
similarity index 57%
rename from pkg/scheduler/tests/mock_rm_callback.go
rename to pkg/mock/rm_callback.go
index 6545cc99..ccaab4a2 100644
--- a/pkg/scheduler/tests/mock_rm_callback.go
+++ b/pkg/mock/rm_callback.go
@@ -16,32 +16,32 @@
limitations under the License.
*/
-package tests
+package mock
import (
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type MockResourceManagerCallback struct{}
+type ResourceManagerCallback struct{}
-func (f *MockResourceManagerCallback) UpdateApplication(response
*si.ApplicationResponse) error {
+func (f *ResourceManagerCallback) UpdateApplication(_ *si.ApplicationResponse)
error {
return nil
}
-func (f *MockResourceManagerCallback) UpdateAllocation(response
*si.AllocationResponse) error {
+func (f *ResourceManagerCallback) UpdateAllocation(_ *si.AllocationResponse)
error {
return nil
}
-func (f *MockResourceManagerCallback) UpdateNode(response *si.NodeResponse)
error {
+func (f *ResourceManagerCallback) UpdateNode(_ *si.NodeResponse) error {
return nil
}
-func (f *MockResourceManagerCallback) Predicates(args *si.PredicatesArgs)
error {
+func (f *ResourceManagerCallback) Predicates(_ *si.PredicatesArgs) error {
// do nothing
return nil
}
-func (f *MockResourceManagerCallback) PreemptionPredicates(args
*si.PreemptionPredicatesArgs) *si.PreemptionPredicatesResponse {
+func (f *ResourceManagerCallback) PreemptionPredicates(args
*si.PreemptionPredicatesArgs) *si.PreemptionPredicatesResponse {
// simulate "ideal" preemption check
return &si.PreemptionPredicatesResponse{
Success: true,
@@ -49,13 +49,14 @@ func (f *MockResourceManagerCallback)
PreemptionPredicates(args *si.PreemptionPr
}
}
-func (f *MockResourceManagerCallback) SendEvent(events []*si.EventRecord) {
+func (f *ResourceManagerCallback) SendEvent(_ []*si.EventRecord) {
// do nothing
}
-func (f *MockResourceManagerCallback) GetStateDump() (string, error) {
+func (f *ResourceManagerCallback) GetStateDump() (string, error) {
return "{}", nil
}
-func (f *MockResourceManagerCallback) UpdateContainerSchedulingState(request
*si.UpdateContainerSchedulingStateRequest) {
+func (f *ResourceManagerCallback) UpdateContainerSchedulingState(_
*si.UpdateContainerSchedulingStateRequest) {
+ // do nothing
}
diff --git a/pkg/plugins/plugins_test.go b/pkg/plugins/plugins_test.go
index 9c5fef12..23b07c38 100644
--- a/pkg/plugins/plugins_test.go
+++ b/pkg/plugins/plugins_test.go
@@ -23,15 +23,11 @@ import (
"gotest.tools/v3/assert"
- "github.com/apache/yunikorn-core/pkg/scheduler/tests"
+ "github.com/apache/yunikorn-core/pkg/mock"
)
-type FakeResourceManagerCallback struct {
- tests.MockResourceManagerCallback
-}
-
func TestRegisterPlugins(t *testing.T) {
plugins = SchedulerPlugins{}
- RegisterSchedulerPlugin(&FakeResourceManagerCallback{})
+ RegisterSchedulerPlugin(&mock.ResourceManagerCallback{})
assert.Assert(t, GetResourceManagerCallbackPlugin() != nil,
"ResourceManagerCallbackPlugin plugin should have been registered")
}
diff --git a/pkg/scheduler/objects/preemption_test.go
b/pkg/scheduler/objects/preemption_test.go
index da8ed670..2b0bf7aa 100644
--- a/pkg/scheduler/objects/preemption_test.go
+++ b/pkg/scheduler/objects/preemption_test.go
@@ -19,17 +19,14 @@
package objects
import (
- "errors"
- "fmt"
"testing"
"time"
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common/resources"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-core/pkg/plugins"
- "github.com/apache/yunikorn-scheduler-interface/lib/go/api"
- "github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
func TestCheckPreconditions(t *testing.T) {
@@ -214,25 +211,15 @@ func TestTryPreemption(t *testing.T) {
preemptor := NewPreemptor(app2, headRoom, 30*time.Second, ask3,
iterator(), false)
// register predicate handler
- pluginErr := &errHolder{}
- plugin := &mockPredicates{
- errHolder: pluginErr,
- reservations: nil,
- allocations: nil,
- preemptions: []mockPreemption{{
- expectedAllocationKey: "alloc3",
- expectedNodeID: "node1",
- expectedAllocationKeys: []string{"alloc1"},
- expectedStartIndex: 0,
- success: true,
- index: 0,
- }},
+ preemptions := []mock.Preemption{
+ mock.NewPreemption(true, "alloc3", "node1", []string{"alloc1"},
0, 0),
}
+ plugin := mock.NewPreemptionPredicatePlugin(nil, nil, preemptions)
plugins.RegisterSchedulerPlugin(plugin)
defer plugins.UnregisterSchedulerPlugins()
alloc, ok := preemptor.TryPreemption()
- assert.NilError(t, pluginErr.err)
+ assert.NilError(t, plugin.GetPredicateError())
assert.Assert(t, ok, "no victims found")
assert.Equal(t, "alloc3", alloc.allocationKey, "wrong alloc")
assert.Check(t, alloc1.IsPreempted(), "alloc1 preempted")
@@ -290,108 +277,3 @@ func allocForScore(originator bool, allowPreemptSelf
bool) *Allocation {
ask.allowPreemptSelf = allowPreemptSelf
return NewAllocation(nodeID1, ask)
}
-
-type mockPreemption struct {
- expectedAllocationKey string
- expectedNodeID string
- expectedAllocationKeys []string
- expectedStartIndex int32
- success bool
- index int32
-}
-
-var _ api.ResourceManagerCallback = &mockPredicates{}
-
-type errHolder struct {
- err error
-}
-
-type mockPredicates struct {
- reservations map[string]string
- allocations map[string]string
- preemptions []mockPreemption
- errHolder *errHolder
- api.ResourceManagerCallback
-}
-
-func (m mockPredicates) Predicates(args *si.PredicatesArgs) error {
- if args.Allocate {
- nodeID, ok := m.allocations[args.AllocationKey]
- if !ok {
- return errors.New("no allocation found")
- }
- if nodeID != args.NodeID {
- return errors.New("wrong node")
- }
- return nil
- } else {
- nodeID, ok := m.reservations[args.AllocationKey]
- if !ok {
- return errors.New("no allocation found")
- }
- if nodeID != args.NodeID {
- return errors.New("wrong node")
- }
- return nil
- }
-}
-
-func (m mockPredicates) PreemptionPredicates(args
*si.PreemptionPredicatesArgs) *si.PreemptionPredicatesResponse {
- result := &si.PreemptionPredicatesResponse{
- Success: false,
- Index: -1,
- }
- for _, preemption := range m.preemptions {
- if preemption.expectedAllocationKey != args.AllocationKey {
- continue
- }
- if preemption.expectedNodeID != args.NodeID {
- continue
- }
- if preemption.expectedStartIndex != args.StartIndex {
- m.errHolder.err = fmt.Errorf("unexpected start index
exepected=%d, actual=%d, allocationKey=%s",
- preemption.expectedStartIndex, args.StartIndex,
args.AllocationKey)
- return result
- }
- if len(preemption.expectedAllocationKeys) !=
len(args.PreemptAllocationKeys) {
- m.errHolder.err = fmt.Errorf("unexpected alloc key
length expected=%d, actual=%d, allocationKey=%s",
- len(preemption.expectedAllocationKeys),
len(args.PreemptAllocationKeys), args.AllocationKey)
- return result
- }
- for idx, key := range preemption.expectedAllocationKeys {
- if args.PreemptAllocationKeys[idx] != key {
- m.errHolder.err = fmt.Errorf("unexpected
preempt alloc key expected=%s, actual=%s, index=%d, allocationKey=%s",
- args.PreemptAllocationKeys[idx], key,
idx, args.AllocationKey)
- return result
- }
- }
- result.Success = preemption.success
- result.Index = preemption.index
- return result
- }
- m.errHolder.err = fmt.Errorf("mo match found allocationKey=%s,
nodeID=%s", args.AllocationKey, args.NodeID)
- return result
-}
-
-func (m mockPredicates) UpdateAllocation(_ *si.AllocationResponse) error {
- // no implementation
- return nil
-}
-
-func (m mockPredicates) UpdateApplication(_ *si.ApplicationResponse) error {
- // no implementation
- return nil
-}
-
-func (m mockPredicates) UpdateNode(_ *si.NodeResponse) error {
- // no implementation
- return nil
-}
-
-func (m mockPredicates) SendEvent(_ []*si.EventRecord) {
- // no implementation
-}
-
-func (m mockPredicates) UpdateContainerSchedulingState(_
*si.UpdateContainerSchedulingStateRequest) {
- // no implementation
-}
diff --git a/pkg/scheduler/partition_test.go b/pkg/scheduler/partition_test.go
index 9237d885..1d605db3 100644
--- a/pkg/scheduler/partition_test.go
+++ b/pkg/scheduler/partition_test.go
@@ -32,6 +32,7 @@ import (
"github.com/apache/yunikorn-core/pkg/common/resources"
"github.com/apache/yunikorn-core/pkg/common/security"
"github.com/apache/yunikorn-core/pkg/events"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-core/pkg/plugins"
"github.com/apache/yunikorn-core/pkg/rmproxy/rmevent"
"github.com/apache/yunikorn-core/pkg/scheduler/objects"
@@ -3344,10 +3345,10 @@ func TestFailReplacePlaceholder(t *testing.T) {
t.Fatalf("empty cluster placeholder allocate returned
allocation: %s", alloc)
}
// plugin to let the pre-check fail on node-1 only, means we cannot
replace the placeholder
- plugin := newFakePredicatePlugin(false, map[string]int{nodeID1: -1})
+ plugin := mock.NewPredicatePlugin(false, map[string]int{nodeID1: -1})
plugins.RegisterSchedulerPlugin(plugin)
defer func() {
- passPlugin := newFakePredicatePlugin(false, nil)
+ passPlugin := mock.NewPredicatePlugin(false, nil)
plugins.RegisterSchedulerPlugin(passPlugin)
}()
var tgRes, res *resources.Resource
diff --git a/pkg/scheduler/tests/mock_rm_callback_test.go
b/pkg/scheduler/tests/mock_rm_callback_test.go
index 53488777..3af34be9 100644
--- a/pkg/scheduler/tests/mock_rm_callback_test.go
+++ b/pkg/scheduler/tests/mock_rm_callback_test.go
@@ -26,11 +26,12 @@ import (
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
type mockRMCallback struct {
- MockResourceManagerCallback
+ mock.ResourceManagerCallback
acceptedApplications map[string]bool
rejectedApplications map[string]bool
acceptedNodes map[string]bool
@@ -41,8 +42,6 @@ type mockRMCallback struct {
sync.RWMutex
}
-// This is only exported to allow the use in the simple_example.go.
-// Lint exclusion added as the non-export return is OK
func newMockRMCallbackHandler() *mockRMCallback {
return &mockRMCallback{
acceptedApplications: make(map[string]bool),
diff --git a/pkg/scheduler/tests/plugin_test.go
b/pkg/scheduler/tests/plugin_test.go
index 3b9f5de6..c45421eb 100644
--- a/pkg/scheduler/tests/plugin_test.go
+++ b/pkg/scheduler/tests/plugin_test.go
@@ -19,35 +19,17 @@
package tests
import (
- "sync"
"testing"
"time"
"gotest.tools/v3/assert"
"github.com/apache/yunikorn-core/pkg/common"
+ "github.com/apache/yunikorn-core/pkg/mock"
"github.com/apache/yunikorn-core/pkg/plugins"
"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
)
-type fakeContainerStateUpdater struct {
- MockResourceManagerCallback
- sentUpdate *si.UpdateContainerSchedulingStateRequest
- sync.RWMutex
-}
-
-func (f *fakeContainerStateUpdater) UpdateContainerSchedulingState(request
*si.UpdateContainerSchedulingStateRequest) {
- f.Lock()
- defer f.Unlock()
- f.sentUpdate = request
-}
-
-func (f *fakeContainerStateUpdater) getContainerUpdateRequest()
*si.UpdateContainerSchedulingStateRequest {
- f.RLock()
- defer f.RUnlock()
- return f.sentUpdate
-}
-
func TestContainerStateUpdater(t *testing.T) {
configData := `
partitions:
@@ -68,8 +50,8 @@ partitions:
assert.NilError(t, err, "RegisterResourceManager failed")
// register a fake container state updater for testing
- fk := &fakeContainerStateUpdater{}
- plugins.RegisterSchedulerPlugin(fk)
+ csu := mock.NewContainerStateUpdater()
+ plugins.RegisterSchedulerPlugin(csu)
const leafName = "root.singleleaf"
const node1 = "node-1"
@@ -149,7 +131,7 @@ partitions:
assert.NilError(t, err)
err = common.WaitFor(100*time.Millisecond, 3000*time.Millisecond,
func() bool {
- reqSent := fk.getContainerUpdateRequest()
+ reqSent := csu.GetContainerUpdateRequest()
return reqSent != nil && reqSent.ApplicationID == appID1 &&
reqSent.GetState() ==
si.UpdateContainerSchedulingStateRequest_FAILED
})
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
