Allow update runner group for AutoScalingRunnerSet (#2216)

github/actions/client.go

@@ -36,6 +36,7 @@ type ActionsService interface {
 	GetRunnerScaleSetById(ctx context.Context, runnerScaleSetId int) (*RunnerScaleSet, error)
 	GetRunnerGroupByName(ctx context.Context, runnerGroup string) (*RunnerGroup, error)
 	CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error)
+	UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error)
 
 	CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*RunnerScaleSetSession, error)
 	DeleteMessageSession(ctx context.Context, runnerScaleSetId int, sessionId *uuid.UUID) error
@@ -350,6 +351,47 @@ func (c *Client) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *Runne
 	return createdRunnerScaleSet, nil
 }
 
+func (c *Client) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) {
+	u := fmt.Sprintf("%s/%s/%d?api-version=6.0-preview", *c.ActionsServiceURL, scaleSetEndpoint, runnerScaleSetId)
+
+	if err := c.refreshTokenIfNeeded(ctx); err != nil {
+		return nil, fmt.Errorf("failed to refresh admin token if needed: %w", err)
+	}
+
+	body, err := json.Marshal(runnerScaleSet)
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPatch, u, bytes.NewBuffer(body))
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *c.ActionsServiceAdminToken))
+
+	if c.userAgent != "" {
+		req.Header.Set("User-Agent", c.userAgent)
+	}
+
+	resp, err := c.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, ParseActionsErrorFromResponse(resp)
+	}
+
+	var updatedRunnerScaleSet *RunnerScaleSet
+	err = unmarshalBody(resp, &updatedRunnerScaleSet)
+	if err != nil {
+		return nil, err
+	}
+	return updatedRunnerScaleSet, nil
+}
+
 func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetId int) error {
 	u := fmt.Sprintf("%s/%s/%d?api-version=6.0-preview", *c.ActionsServiceURL, scaleSetEndpoint, runnerScaleSetId)
 

github/actions/client_runner_scale_set_test.go

@@ -336,3 +336,53 @@ func TestCreateRunnerScaleSet(t *testing.T) {
 		assert.Equalf(t, actualRetry, expectedRetry, "A retry was expected after the first request but got: %v", actualRetry)
 	})
 }
+
+func TestUpdateRunnerScaleSet(t *testing.T) {
+	ctx := context.Background()
+	auth := &actions.ActionsAuth{
+		Token: "token",
+	}
+
+	scaleSetCreationDateTime := time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC)
+	runnerScaleSet := actions.RunnerScaleSet{Id: 1, Name: "ScaleSet", RunnerGroupId: 1, RunnerGroupName: "group", CreatedOn: scaleSetCreationDateTime, RunnerSetting: actions.RunnerSetting{}}
+
+	t.Run("Update runner scale set", func(t *testing.T) {
+		want := &runnerScaleSet
+		rsl, err := json.Marshal(want)
+		require.NoError(t, err)
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Write(rsl)
+		}))
+
+		client, err := actions.NewClient(ctx, server.configURLForOrg("my-org"), auth)
+		require.NoError(t, err)
+
+		got, err := client.UpdateRunnerScaleSet(ctx, 1, &actions.RunnerScaleSet{RunnerGroupId: 1})
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("UpdateRunnerScaleSet calls correct url", func(t *testing.T) {
+		rsl, err := json.Marshal(&runnerScaleSet)
+		require.NoError(t, err)
+		url := url.URL{}
+		method := ""
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write(rsl)
+			url = *r.URL
+			method = r.Method
+		}))
+
+		client, err := actions.NewClient(ctx, server.configURLForOrg("my-org"), auth)
+		require.NoError(t, err)
+
+		_, err = client.UpdateRunnerScaleSet(ctx, 1, &runnerScaleSet)
+		require.NoError(t, err)
+
+		u := url.String()
+		expectedUrl := "/_apis/runtime/runnerscalesets/1?api-version=6.0-preview"
+		assert.Equal(t, expectedUrl, u)
+
+		assert.Equal(t, "PATCH", method)
+	})
+}

github/actions/fake/client.go

@@ -36,6 +36,18 @@ var defaultRunnerScaleSet = &actions.RunnerScaleSet{
 	Statistics:         nil,
 }
 
+var defaultUpdatedRunnerScaleSet = &actions.RunnerScaleSet{
+	Id:                 1,
+	Name:               "testset",
+	RunnerGroupId:      2,
+	RunnerGroupName:    "testgroup",
+	Labels:             []actions.Label{{Type: "test", Name: "test"}},
+	RunnerSetting:      actions.RunnerSetting{},
+	CreatedOn:          time.Now(),
+	RunnerJitConfigUrl: "test.test.test",
+	Statistics:         nil,
+}
+
 var defaultRunnerGroup = &actions.RunnerGroup{
 	ID:        1,
 	Name:      "testgroup",
@@ -107,6 +119,10 @@ type FakeClient struct {
 		*actions.RunnerScaleSet
 		err error
 	}
+	updateRunnerScaleSetResult struct {
+		*actions.RunnerScaleSet
+		err error
+	}
 	createMessageSessionResult struct {
 		*actions.RunnerScaleSetSession
 		err error
@@ -164,6 +180,7 @@ func (f *FakeClient) applyDefaults() {
 	f.getRunnerScaleSetByIdResult.RunnerScaleSet = defaultRunnerScaleSet
 	f.getRunnerGroupByNameResult.RunnerGroup = defaultRunnerGroup
 	f.createRunnerScaleSetResult.RunnerScaleSet = defaultRunnerScaleSet
+	f.updateRunnerScaleSetResult.RunnerScaleSet = defaultUpdatedRunnerScaleSet
 	f.createMessageSessionResult.RunnerScaleSetSession = defaultRunnerScaleSetSession
 	f.refreshMessageSessionResult.RunnerScaleSetSession = defaultRunnerScaleSetSession
 	f.acquireJobsResult.ids = []int64{1}
@@ -190,6 +207,10 @@ func (f *FakeClient) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *a
 	return f.createRunnerScaleSetResult.RunnerScaleSet, f.createRunnerScaleSetResult.err
 }
 
+func (f *FakeClient) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *actions.RunnerScaleSet) (*actions.RunnerScaleSet, error) {
+	return f.updateRunnerScaleSetResult.RunnerScaleSet, f.updateRunnerScaleSetResult.err
+}
+
 func (f *FakeClient) CreateMessageSession(ctx context.Context, runnerScaleSetId int, owner string) (*actions.RunnerScaleSetSession, error) {
 	return f.createMessageSessionResult.RunnerScaleSetSession, f.createMessageSessionResult.err
 }

github/actions/mock_ActionsService.go

@@ -332,6 +332,29 @@ func (_m *MockActionsService) RemoveRunner(ctx context.Context, runnerId int64)
 	return r0
 }
 
+// UpdateRunnerScaleSet provides a mock function with given fields: ctx, runnerScaleSetId, runnerScaleSet
+func (_m *MockActionsService) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetId int, runnerScaleSet *RunnerScaleSet) (*RunnerScaleSet, error) {
+	ret := _m.Called(ctx, runnerScaleSetId, runnerScaleSet)
+
+	var r0 *RunnerScaleSet
+	if rf, ok := ret.Get(0).(func(context.Context, int, *RunnerScaleSet) *RunnerScaleSet); ok {
+		r0 = rf(ctx, runnerScaleSetId, runnerScaleSet)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*RunnerScaleSet)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, int, *RunnerScaleSet) error); ok {
+		r1 = rf(ctx, runnerScaleSetId, runnerScaleSet)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 type mockConstructorTestingTNewMockActionsService interface {
 	mock.TestingT
 	Cleanup(func())