diff --git a/realis.go b/realis.go index a1512c7..dde0e6d 100644 --- a/realis.go +++ b/realis.go @@ -71,7 +71,7 @@ type Realis interface { DrainHosts(hosts ...string) (*aurora.Response, *aurora.DrainHostsResult_, error) EndMaintenance(hosts ...string) (*aurora.Response, *aurora.EndMaintenanceResult_, error) MaintenanceStatus(hosts ...string) (*aurora.Response, *aurora.MaintenanceStatusResult_, error) - SetQuota(role string, cpu float64, ram int64, disk int64) (*aurora.Response, error) + SetQuota(role string, cpu *float64, ram *int64, disk *int64) (*aurora.Response, error) GetQuota(role string) (*aurora.Response, error) } @@ -948,20 +948,28 @@ func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *au } // SetQuota sets a quota aggregate for the given role -func (r *realisClient) SetQuota(role string, cpu float64, ramMb int64, diskMb int64) (*aurora.Response, error) { - quota := &aurora.ResourceAggregate{ - NumCpus: cpu, - RamMb: ramMb, - DiskMb: diskMb, - } +// TODO(zircote) Currently investigating an error that is returned from thrift calls that include resources for `NamedPort` and `NumGpu` +func (r *realisClient) SetQuota(role string, cpu *float64, ramMb *int64, diskMb *int64) (*aurora.Response, error) { + ram := aurora.NewResource() + ram.RamMb = ramMb + c := aurora.NewResource() + c.NumCpus = cpu + d := aurora.NewResource() + d.DiskMb = diskMb + quota := aurora.NewResourceAggregate() + quota.Resources = make(map[*aurora.Resource]bool) + quota.Resources[ram] = true + quota.Resources[c] = true + quota.Resources[d] = true resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { - return r.adminClient.SetQuota(role, quota) - }) + resp, retryErr := r.adminClient.SetQuota(role, quota) - if retryErr != nil { - return nil, errors.Wrap(retryErr, "Unable to set role quota") - } - return resp, nil + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to set role quota") + } + return resp, nil + }) + return resp, retryErr } @@ -969,11 +977,13 @@ func (r *realisClient) SetQuota(role string, cpu float64, ramMb int64, diskMb in func (r *realisClient) GetQuota(role string) (*aurora.Response, error) { resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { - return r.adminClient.GetQuota(role) + resp, retryErr :=r.adminClient.GetQuota(role) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to get role quota") + } + return resp, nil }) - if retryErr != nil { - return nil, errors.Wrap(retryErr, "Unable to get role quota") - } - return resp, nil + return resp, retryErr } diff --git a/realis_e2e_test.go b/realis_e2e_test.go index 1528594..ae61cf7 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -392,23 +392,34 @@ func TestRealisClient_SessionThreadSafety(t *testing.T) { // Test setting and getting the quota func TestRealisClient_SetQuota(t *testing.T) { - var cpu = 2.5 - var ram int64 = 10240 + var cpu = 3.5 + var ram int64 = 20480 var disk int64 = 10240 - resp, err := r.SetQuota("vagrant", cpu, ram, disk) + resp, err := r.SetQuota("vagrant", &cpu, &ram, &disk) assert.NoError(t, err) assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) - - // Test GetQuota based on previously set values - var result *aurora.GetQuotaResult_ - resp, err = r.GetQuota("vagrant") - if resp.GetResult_() != nil { - result = resp.GetResult_().GetQuotaResult_ - } - assert.NoError(t, err) - assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) - assert.Equal(t, cpu, result.Quota.NumCpus) - assert.Equal(t, ram, result.Quota.RamMb) - assert.Equal(t, disk, result.Quota.DiskMb) - fmt.Print("GetQuota Result", result.String()) + t.Run("TestRealisClient_GetQuota", func(t *testing.T) { + // Test GetQuota based on previously set values + var result *aurora.GetQuotaResult_ + resp, err = r.GetQuota("vagrant") + if resp.GetResult_() != nil { + result = resp.GetResult_().GetQuotaResult_ + } + assert.NoError(t, err) + assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) + for res := range result.Quota.GetResources() { + switch true { + case res.DiskMb != nil: + assert.Equal(t, disk, *res.DiskMb) + break + case res.NumCpus != nil: + assert.Equal(t, cpu, *res.NumCpus) + break + case res.RamMb != nil: + assert.Equal(t, ram, *res.RamMb) + break + } + } + fmt.Print("GetQuota Result", result.String()) + }) }