Adding Admin Client calls GetQuota & SetQuota (#59)

* Adding Admin Client calls `GetQuota` & `SetQuota`

This change set adds admin client calls to fetch and
mutate the OwnerRole quota[cpu,ram,disk].

realis.go

@@ -71,6 +71,8 @@ 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)
+	GetQuota(role string) (*aurora.Response, error)
 }
 
 type realisClient struct {
@@ -944,3 +946,44 @@ func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *au
 
 	return resp, result, nil
 }
+
+// SetQuota sets a quota aggregate for the given role
+// 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) {
+		resp, retryErr := r.adminClient.SetQuota(role, quota)
+
+		if retryErr != nil {
+			return nil, errors.Wrap(retryErr, "Unable to set role quota")
+		}
+		return resp, nil
+	})
+	return resp, retryErr
+
+}
+
+// GetQuota returns the resource aggregate for the given role
+func (r *realisClient) GetQuota(role string) (*aurora.Response, error) {
+
+	resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
+		resp, retryErr :=r.adminClient.GetQuota(role)
+
+		if retryErr != nil {
+			return nil, errors.Wrap(retryErr, "Unable to get role quota")
+		}
+		return resp, nil
+	})
+
+	return resp, retryErr
+}

realis_e2e_test.go

@@ -389,3 +389,37 @@ func TestRealisClient_SessionThreadSafety(t *testing.T) {
 
 	wg.Wait()
 }
+
+// Test setting and getting the quota
+func TestRealisClient_SetQuota(t *testing.T) {
+	var cpu = 3.5
+	var ram int64 = 20480
+	var disk int64 = 10240
+	resp, err := r.SetQuota("vagrant", &cpu, &ram, &disk)
+	assert.NoError(t, err)
+	assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode)
+	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())
+	})
+}

updatejob.go

@@ -138,7 +138,6 @@ func (u *UpdateJob) RollbackOnFail(rollback bool) *UpdateJob {
 	return u
 }
 
-
 func NewUpdateSettings() *aurora.JobUpdateSettings {
 
 	us := new(aurora.JobUpdateSettings)