github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/resource/plugins/cpumem/cpumem_test.go (about)

     1  package cpumem
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  
     8  	enginetypes "github.com/projecteru2/core/engine/types"
     9  	plugintypes "github.com/projecteru2/core/resource/plugins/types"
    10  	coretypes "github.com/projecteru2/core/types"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  func TestName(t *testing.T) {
    15  	cm := initCPUMEM(context.Background(), t)
    16  	assert.Equal(t, cm.name, cm.Name())
    17  }
    18  
    19  func initCPUMEM(ctx context.Context, t *testing.T) *Plugin {
    20  	config := coretypes.Config{
    21  		Etcd: coretypes.EtcdConfig{
    22  			Prefix: "/cpumem",
    23  		},
    24  		Scheduler: coretypes.SchedulerConfig{
    25  			MaxShare:  -1,
    26  			ShareBase: 100,
    27  		},
    28  	}
    29  
    30  	cm, err := NewPlugin(ctx, config, t)
    31  	assert.NoError(t, err)
    32  	return cm
    33  }
    34  
    35  func generateNodes(
    36  	ctx context.Context, t *testing.T, cm *Plugin,
    37  	nums int, cores int, memory int64, shares, index int,
    38  ) []string {
    39  	reqs := generateNodeResourceRequests(t, nums, cores, memory, shares, index)
    40  	info := &enginetypes.Info{NCPU: 8, MemTotal: 2048}
    41  	names := []string{}
    42  	for name, req := range reqs {
    43  		_, err := cm.AddNode(ctx, name, req, info)
    44  		assert.NoError(t, err)
    45  		names = append(names, name)
    46  	}
    47  	t.Cleanup(func() {
    48  		for name := range reqs {
    49  			cm.RemoveNode(ctx, name)
    50  		}
    51  	})
    52  	return names
    53  }
    54  
    55  func generateNodeResourceRequests(t *testing.T, nums int, cores int, memory int64, shares, index int) map[string]plugintypes.NodeResourceRequest {
    56  	infos := map[string]plugintypes.NodeResourceRequest{}
    57  	for i := index; i < index+nums; i++ {
    58  		info := plugintypes.NodeResourceRequest{
    59  			"cpu":    cores,
    60  			"share":  shares,
    61  			"memory": fmt.Sprintf("%v", memory),
    62  		}
    63  		infos[fmt.Sprintf("test%v", i)] = info
    64  	}
    65  	return infos
    66  }