From c2309468225be00260a83f1d8724fca4031b3e0a Mon Sep 17 00:00:00 2001 From: Anuj Singh Date: Wed, 11 Sep 2024 14:01:10 -0700 Subject: [PATCH] update unit tests --- agent/engine/host_resource_manager.go | 26 +++--- agent/engine/host_resource_manager_test.go | 101 ++++++++++++--------- 2 files changed, 72 insertions(+), 55 deletions(-) diff --git a/agent/engine/host_resource_manager.go b/agent/engine/host_resource_manager.go index 602c73cddbc..d505f96bfde 100644 --- a/agent/engine/host_resource_manager.go +++ b/agent/engine/host_resource_manager.go @@ -115,7 +115,7 @@ func (h *HostResourceManager) consume(taskArn string, resources map[string]*ecs. return true, nil } - ok, err, failedResourceKey := h.consumable(resources) + ok, failedResourceKey, err := h.consumable(resources) if err != nil { logger.Error("Resources failing to consume, error in task resources", logger.Fields{ "taskArn": taskArn, @@ -140,7 +140,7 @@ func (h *HostResourceManager) consume(taskArn string, resources map[string]*ecs. return true, nil } logger.Info("Resources not consumed, enough resources not available", logger.Fields{ - "taskArn": taskArn, + "taskArn": taskArn, "resource": failedResourceKey, }) return false, nil @@ -256,29 +256,29 @@ func (h *HostResourceManager) checkResourcesHealth(resources map[string]*ecs.Res // Helper function for consume to check if resources are consumable with the current account // we have for the host resources. Should not call host resource manager lock in this func // return values -func (h *HostResourceManager) consumable(resources map[string]*ecs.Resource) (bool, error, string) { +func (h *HostResourceManager) consumable(resources map[string]*ecs.Resource) (bool, string, error) { err := h.checkResourcesHealth(resources) if err != nil { - return false, err, nil + return false, "", err } for resourceKey := range resources { if *resources[resourceKey].Type == "INTEGER" { consumable := h.checkConsumableIntType(resourceKey, resources) if !consumable { - return false, nil, resourceKey + return false, resourceKey, nil } } if *resources[resourceKey].Type == "STRINGSET" { consumable := h.checkConsumableStringSetType(resourceKey, resources) if !consumable { - return false, nil, resourceKey + return false, resourceKey, nil } } } - return true, nil, nil + return true, "", nil } // Utility function to manage release of ports @@ -351,22 +351,22 @@ func NewHostResourceManager(resourceMap map[string]*ecs.Resource) HostResourceMa consumedResourceMap := make(map[string]*ecs.Resource) taskConsumed := make(map[string]bool) // assigns CPU, MEMORY, PORTS_TCP, PORTS_UDP from host - //CPU + // CPU CPUs := int64(0) consumedResourceMap[CPU] = &ecs.Resource{ Name: utils.Strptr(CPU), Type: utils.Strptr("INTEGER"), IntegerValue: &CPUs, } - //MEMORY + // MEMORY memory := int64(0) consumedResourceMap[MEMORY] = &ecs.Resource{ Name: utils.Strptr(MEMORY), Type: utils.Strptr("INTEGER"), IntegerValue: &memory, } - //PORTS_TCP - //Copying ports from host resources as consumed ports for initializing + // PORTS_TCP + // Copying ports from host resources as consumed ports for initializing portsTcp := []*string{} if resourceMap != nil && resourceMap[PORTSTCP] != nil { portsTcp = resourceMap[PORTSTCP].StringSetValue @@ -377,7 +377,7 @@ func NewHostResourceManager(resourceMap map[string]*ecs.Resource) HostResourceMa StringSetValue: portsTcp, } - //PORTS_UDP + // PORTS_UDP portsUdp := []*string{} if resourceMap != nil && resourceMap[PORTSUDP] != nil { portsUdp = resourceMap[PORTSUDP].StringSetValue @@ -388,7 +388,7 @@ func NewHostResourceManager(resourceMap map[string]*ecs.Resource) HostResourceMa StringSetValue: portsUdp, } - //GPUs + // GPUs gpuIDs := []*string{} consumedResourceMap[GPU] = &ecs.Resource{ Name: utils.Strptr(GPU), diff --git a/agent/engine/host_resource_manager_test.go b/agent/engine/host_resource_manager_test.go index 21715386333..f04e82a4d24 100644 --- a/agent/engine/host_resource_manager_test.go +++ b/agent/engine/host_resource_manager_test.go @@ -54,7 +54,7 @@ func getTestHostResourceManager(cpu int64, mem int64, ports []*string, portsUdp } hostResources["GPU"] = &ecs.Resource{ - Name: utils.Strptr("PORTS_UDP"), + Name: utils.Strptr("GPU"), Type: utils.Strptr("STRINGSET"), StringSetValue: gpuIDs, } @@ -180,63 +180,80 @@ func TestConsumable(t *testing.T) { hostResourcePort1 := "22" hostResourcePort2 := "1000" gpuIDs := []string{"gpu1", "gpu2", "gpu3", "gpu4"} - h := getTestHostResourceManager(int64(2048), int64(2048), []*string{&hostResourcePort1}, []*string{&hostResourcePort2}, aws.StringSlice(gpuIDs)) + h := getTestHostResourceManager(int64(2048), int64(2048), []*string{&hostResourcePort1}, + []*string{&hostResourcePort2}, aws.StringSlice(gpuIDs)) testCases := []struct { - cpu int64 - mem int64 - ports []uint16 - portsUdp []uint16 - gpus []string - canBeConsumed bool + name string + cpu int64 + mem int64 + ports []uint16 + portsUdp []uint16 + gpus []string + canBeConsumed bool + expectedFailedResourceKey string }{ { - cpu: int64(1024), - mem: int64(1024), - ports: []uint16{25}, - portsUdp: []uint16{1003}, - gpus: []string{"gpu1", "gpu2"}, - canBeConsumed: true, + name: "consumable", + cpu: int64(1024), + mem: int64(1024), + ports: []uint16{25}, + portsUdp: []uint16{1003}, + gpus: []string{"gpu1", "gpu2"}, + canBeConsumed: true, + expectedFailedResourceKey: "", }, { - cpu: int64(2500), - mem: int64(1024), - ports: []uint16{}, - portsUdp: []uint16{}, - gpus: []string{}, - canBeConsumed: false, + name: "cpu not consumable", + cpu: int64(2500), + mem: int64(1024), + ports: []uint16{}, + portsUdp: []uint16{}, + gpus: []string{}, + canBeConsumed: false, + expectedFailedResourceKey: "CPU", }, { - cpu: int64(1024), - mem: int64(2500), - ports: []uint16{}, - portsUdp: []uint16{}, - gpus: []string{}, - canBeConsumed: false, + name: "memory not consumable", + cpu: int64(1024), + mem: int64(2500), + ports: []uint16{}, + portsUdp: []uint16{}, + gpus: []string{}, + canBeConsumed: false, + expectedFailedResourceKey: "MEMORY", }, { - cpu: int64(1024), - mem: int64(1024), - ports: []uint16{22}, - portsUdp: []uint16{}, - gpus: []string{}, - canBeConsumed: false, + name: "tcp ports not consumable", + cpu: int64(1024), + mem: int64(1024), + ports: []uint16{22}, + portsUdp: []uint16{}, + gpus: []string{}, + canBeConsumed: false, + expectedFailedResourceKey: "PORTS_TCP", }, { - cpu: int64(1024), - mem: int64(1024), - ports: []uint16{}, - portsUdp: []uint16{1000}, - gpus: []string{}, - canBeConsumed: false, + name: "udp ports not consumable", + cpu: int64(1024), + mem: int64(1024), + ports: []uint16{}, + portsUdp: []uint16{1000}, + gpus: []string{}, + canBeConsumed: false, + expectedFailedResourceKey: "PORTS_UDP", }, } for _, tc := range testCases { - resources := getTestTaskResourceMap(tc.cpu, tc.mem, commonutils.Uint16SliceToStringSlice(tc.ports), commonutils.Uint16SliceToStringSlice(tc.portsUdp), aws.StringSlice(tc.gpus)) - canBeConsumed, err := h.consumable(resources) - assert.Equal(t, canBeConsumed, tc.canBeConsumed, "Error in checking if resources can be successfully consumed") - assert.Equal(t, err, nil, "Error in checking if resources can be successfully consumed, error returned from consumable") + resources := getTestTaskResourceMap(tc.cpu, tc.mem, commonutils.Uint16SliceToStringSlice(tc.ports), + commonutils.Uint16SliceToStringSlice(tc.portsUdp), aws.StringSlice(tc.gpus)) + canBeConsumed, failedResourceKey, err := h.consumable(resources) + assert.Equal(t, tc.canBeConsumed, canBeConsumed, + "Error in checking if resources can be successfully consumed") + assert.Equal(t, nil, err, + "Error in checking if resources can be successfully consumed, error returned from consumable") + assert.Equal(t, tc.expectedFailedResourceKey, failedResourceKey) } }