Moving to scaleset client for the controller (#4390)

2026-03-13 14:36:41 +01:00
parent 1d9f626c53
commit f99c6eda0b
27 changed files with 2695 additions and 360 deletions
@@ -0,0 +1,230 @@
+package fake
+
+import (
+	"context"
+
+	"github.com/actions/actions-runner-controller/controllers/actions.github.com/multiclient"
+	"github.com/actions/scaleset"
+)
+
+// ClientOption is a functional option for configuring a fake Client
+type ClientOption func(*Client)
+
+// WithGetRunnerScaleSet configures the result of GetRunnerScaleSet
+func WithGetRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerScaleSetResult.RunnerScaleSet = result
+		c.getRunnerScaleSetResult.err = err
+	}
+}
+
+// WithGetRunnerScaleSetByID configures the result of GetRunnerScaleSetByID
+func WithGetRunnerScaleSetByID(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerScaleSetByIDResult.RunnerScaleSet = result
+		c.getRunnerScaleSetByIDResult.err = err
+	}
+}
+
+// WithGetRunnerGroupByName configures the result of GetRunnerGroupByName
+func WithGetRunnerGroupByName(result *scaleset.RunnerGroup, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerGroupByNameResult.RunnerGroup = result
+		c.getRunnerGroupByNameResult.err = err
+	}
+}
+
+// WithGetRunnerGroupByNameFunc configures a function to handle GetRunnerGroupByName calls dynamically
+func WithGetRunnerGroupByNameFunc(fn func(context.Context, string) (*scaleset.RunnerGroup, error)) ClientOption {
+	return func(c *Client) {
+		c.getRunnerGroupByNameFunc = fn
+	}
+}
+
+// WithCreateRunnerScaleSet configures the result of CreateRunnerScaleSet
+func WithCreateRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.createRunnerScaleSetResult.RunnerScaleSet = result
+		c.createRunnerScaleSetResult.err = err
+	}
+}
+
+// WithUpdateRunnerScaleSet configures the result of UpdateRunnerScaleSet
+func WithUpdateRunnerScaleSet(result *scaleset.RunnerScaleSet, err error) ClientOption {
+	return func(c *Client) {
+		c.updateRunnerScaleSetResult.RunnerScaleSet = result
+		c.updateRunnerScaleSetResult.err = err
+	}
+}
+
+// WithDeleteRunnerScaleSet configures the result of DeleteRunnerScaleSet
+func WithDeleteRunnerScaleSet(err error) ClientOption {
+	return func(c *Client) {
+		c.deleteRunnerScaleSetResult.err = err
+	}
+}
+
+// WithRemoveRunner configures the result of RemoveRunner
+func WithRemoveRunner(err error) ClientOption {
+	return func(c *Client) {
+		c.removeRunnerResult.err = err
+	}
+}
+
+// WithGenerateJitRunnerConfig configures the result of GenerateJitRunnerConfig
+func WithGenerateJitRunnerConfig(result *scaleset.RunnerScaleSetJitRunnerConfig, err error) ClientOption {
+	return func(c *Client) {
+		c.generateJitRunnerConfigResult.RunnerScaleSetJitRunnerConfig = result
+		c.generateJitRunnerConfigResult.err = err
+	}
+}
+
+// WithGetRunnerByName configures the result of GetRunnerByName
+func WithGetRunnerByName(result *scaleset.RunnerReference, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerByNameResult.RunnerReference = result
+		c.getRunnerByNameResult.err = err
+	}
+}
+
+// WithGetRunner configures the result of GetRunner
+func WithGetRunner(result *scaleset.RunnerReference, err error) ClientOption {
+	return func(c *Client) {
+		c.getRunnerResult.RunnerReference = result
+		c.getRunnerResult.err = err
+	}
+}
+
+// WithSystemInfo configures the SystemInfo
+func WithSystemInfo(info scaleset.SystemInfo) ClientOption {
+	return func(c *Client) {
+		c.systemInfo = info
+	}
+}
+
+// WithUpdateRunnerScaleSetFunc configures a function to handle UpdateRunnerScaleSet calls dynamically
+func WithUpdateRunnerScaleSetFunc(fn func(context.Context, int, *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)) ClientOption {
+	return func(c *Client) {
+		c.updateRunnerScaleSetFunc = fn
+	}
+}
+
+// Client implements multiclient.Client interface for testing
+type Client struct {
+	systemInfo               scaleset.SystemInfo
+	updateRunnerScaleSetFunc func(context.Context, int, *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error)
+
+	getRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	getRunnerScaleSetByIDResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	getRunnerGroupByNameResult struct {
+		*scaleset.RunnerGroup
+		err error
+	}
+	getRunnerGroupByNameFunc   func(context.Context, string) (*scaleset.RunnerGroup, error)
+	createRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	updateRunnerScaleSetResult struct {
+		*scaleset.RunnerScaleSet
+		err error
+	}
+	deleteRunnerScaleSetResult struct {
+		err error
+	}
+	removeRunnerResult struct {
+		err error
+	}
+	generateJitRunnerConfigResult struct {
+		*scaleset.RunnerScaleSetJitRunnerConfig
+		err error
+	}
+	getRunnerByNameResult struct {
+		*scaleset.RunnerReference
+		err error
+	}
+	getRunnerResult struct {
+		*scaleset.RunnerReference
+		err error
+	}
+	messageSessionClientResult struct {
+		*scaleset.MessageSessionClient
+		err error
+	}
+}
+
+// Compile-time interface check
+var _ multiclient.Client = (*Client)(nil)
+
+// NewClient creates a new fake Client with the given options
+func NewClient(opts ...ClientOption) *Client {
+	c := &Client{}
+	for _, opt := range opts {
+		opt(c)
+	}
+	return c
+}
+
+func (c *Client) SetSystemInfo(info scaleset.SystemInfo) {
+	c.systemInfo = info
+}
+
+func (c *Client) SystemInfo() scaleset.SystemInfo {
+	return c.systemInfo
+}
+
+func (c *Client) MessageSessionClient(ctx context.Context, runnerScaleSetID int, owner string, options ...scaleset.HTTPOption) (*scaleset.MessageSessionClient, error) {
+	return c.messageSessionClientResult.MessageSessionClient, c.messageSessionClientResult.err
+}
+
+func (c *Client) GenerateJitRunnerConfig(ctx context.Context, jitRunnerSetting *scaleset.RunnerScaleSetJitRunnerSetting, scaleSetID int) (*scaleset.RunnerScaleSetJitRunnerConfig, error) {
+	return c.generateJitRunnerConfigResult.RunnerScaleSetJitRunnerConfig, c.generateJitRunnerConfigResult.err
+}
+
+func (c *Client) GetRunner(ctx context.Context, runnerID int) (*scaleset.RunnerReference, error) {
+	return c.getRunnerResult.RunnerReference, c.getRunnerResult.err
+}
+
+func (c *Client) GetRunnerByName(ctx context.Context, runnerName string) (*scaleset.RunnerReference, error) {
+	return c.getRunnerByNameResult.RunnerReference, c.getRunnerByNameResult.err
+}
+
+func (c *Client) RemoveRunner(ctx context.Context, runnerID int64) error {
+	return c.removeRunnerResult.err
+}
+
+func (c *Client) GetRunnerGroupByName(ctx context.Context, runnerGroup string) (*scaleset.RunnerGroup, error) {
+	if c.getRunnerGroupByNameFunc != nil {
+		return c.getRunnerGroupByNameFunc(ctx, runnerGroup)
+	}
+	return c.getRunnerGroupByNameResult.RunnerGroup, c.getRunnerGroupByNameResult.err
+}
+
+func (c *Client) GetRunnerScaleSet(ctx context.Context, runnerGroupID int, runnerScaleSetName string) (*scaleset.RunnerScaleSet, error) {
+	return c.getRunnerScaleSetResult.RunnerScaleSet, c.getRunnerScaleSetResult.err
+}
+
+func (c *Client) GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID int) (*scaleset.RunnerScaleSet, error) {
+	return c.getRunnerScaleSetByIDResult.RunnerScaleSet, c.getRunnerScaleSetByIDResult.err
+}
+
+func (c *Client) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error) {
+	return c.createRunnerScaleSetResult.RunnerScaleSet, c.createRunnerScaleSetResult.err
+}
+
+func (c *Client) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetID int, runnerScaleSet *scaleset.RunnerScaleSet) (*scaleset.RunnerScaleSet, error) {
+	if c.updateRunnerScaleSetFunc != nil {
+		return c.updateRunnerScaleSetFunc(ctx, runnerScaleSetID, runnerScaleSet)
+	}
+	return c.updateRunnerScaleSetResult.RunnerScaleSet, c.updateRunnerScaleSetResult.err
+}
+
+func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetID int) error {
+	return c.deleteRunnerScaleSetResult.err
+}
@@ -0,0 +1,53 @@
+package fake
+
+import (
+	"context"
+
+	"github.com/actions/actions-runner-controller/controllers/actions.github.com/multiclient"
+)
+
+// MultiClientOption is a functional option for configuring a fake MultiClient
+type MultiClientOption func(*MultiClient)
+
+// WithClient configures the client that GetClientFor will return
+func WithClient(c multiclient.Client) MultiClientOption {
+	return func(mc *MultiClient) {
+		mc.client = c
+	}
+}
+
+// WithGetClientForError configures an error that GetClientFor will return
+func WithGetClientForError(err error) MultiClientOption {
+	return func(mc *MultiClient) {
+		mc.getClientForErr = err
+	}
+}
+
+// MultiClient implements multiclient.MultiClient interface for testing
+type MultiClient struct {
+	client          multiclient.Client
+	getClientForErr error
+}
+
+// Compile-time interface check
+var _ multiclient.MultiClient = (*MultiClient)(nil)
+
+// NewMultiClient creates a new fake MultiClient with the given options
+func NewMultiClient(opts ...MultiClientOption) *MultiClient {
+	mc := &MultiClient{}
+	for _, opt := range opts {
+		opt(mc)
+	}
+	// Default behavior: if no client configured, return a default NewClient()
+	if mc.client == nil {
+		mc.client = NewClient()
+	}
+	return mc
+}
+
+func (mc *MultiClient) GetClientFor(ctx context.Context, opts *multiclient.ClientForOptions) (multiclient.Client, error) {
+	if mc.getClientForErr != nil {
+		return nil, mc.getClientForErr
+	}
+	return mc.client, nil
+}