Modify error API and add Listener tests (#36)

* Modify error API to not include exception, but to test on type

* reformat

* address renames and documentations from the review

* test is errors

* document

.github/workflows/go.yaml

@@ -29,7 +29,7 @@ jobs:
       - name: Check diff
         run: git diff --exit-code
 
-  generate:
+  mocks:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v5
@@ -37,8 +37,8 @@ jobs:
         with:
           go-version-file: "go.mod"
           cache: false
-      - name: go generate
-        run: go generate ./...
+      - name: "Run mockery"
+        run: go tool github.com/vektra/mockery/v3
       - name: Check diff
         run: git diff --exit-code
 

.mockery.yml

@@ -0,0 +1,14 @@
+all: false
+dir: "{{.InterfaceDir}}"
+filename: mocks_test.go
+force-file-write: true
+formatter: goimports
+log-level: info
+structname: "{{.Mock}}{{.InterfaceName}}"
+pkgname: "{{.SrcPackageName}}"
+recursive: true
+template: testify
+packages:
+  github.com/actions/scaleset/listener:
+    config:
+      all: true

client.go

@@ -361,7 +361,7 @@ func (c *Client) GetRunnerScaleSet(ctx context.Context, runnerGroupID int, runne
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var runnerScaleSetList *runnerScaleSetsResponse
@@ -402,7 +402,7 @@ func (c *Client) GetRunnerScaleSetByID(ctx context.Context, runnerScaleSetID int
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var runnerScaleSet *RunnerScaleSet
@@ -492,7 +492,7 @@ func (c *Client) CreateRunnerScaleSet(ctx context.Context, runnerScaleSet *Runne
 	}
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 	var createdRunnerScaleSet *RunnerScaleSet
 	if err := json.NewDecoder(resp.Body).Decode(&createdRunnerScaleSet); err != nil {
@@ -525,7 +525,7 @@ func (c *Client) UpdateRunnerScaleSet(ctx context.Context, runnerScaleSetID int,
 	}
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var updatedRunnerScaleSet *RunnerScaleSet
@@ -554,7 +554,7 @@ func (c *Client) DeleteRunnerScaleSet(ctx context.Context, runnerScaleSetID int)
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusNoContent {
-		return parseActionsErrorFromResponse(resp)
+		return ParseActionsErrorFromResponse(resp)
 	}
 
 	return nil
@@ -598,7 +598,7 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc
 
 	if resp.StatusCode != http.StatusOK {
 		if resp.StatusCode != http.StatusUnauthorized {
-			return nil, parseActionsErrorFromResponse(resp)
+			return nil, ParseActionsErrorFromResponse(resp)
 		}
 
 		body, err := io.ReadAll(resp.Body)
@@ -610,10 +610,12 @@ func (c *Client) GetMessage(ctx context.Context, messageQueueURL, messageQueueAc
 				Err:        err,
 			}
 		}
-		return nil, &MessageQueueTokenExpiredError{
-			activityID: resp.Header.Get(headerActionsActivityID),
-			statusCode: resp.StatusCode,
-			msg:        string(body),
+		return nil, &ActionsError{
+			ActivityID: resp.Header.Get(headerActionsActivityID),
+			StatusCode: resp.StatusCode,
+			Err: &messageQueueTokenExpiredError{
+				message: string(body),
+			},
 		}
 	}
 
@@ -721,7 +723,7 @@ func (c *Client) DeleteMessage(ctx context.Context, messageQueueURL, messageQueu
 	}
 
 	if resp.StatusCode != http.StatusUnauthorized {
-		return parseActionsErrorFromResponse(resp)
+		return ParseActionsErrorFromResponse(resp)
 	}
 
 	body, err := io.ReadAll(resp.Body)
@@ -733,10 +735,12 @@ func (c *Client) DeleteMessage(ctx context.Context, messageQueueURL, messageQueu
 			Err:        err,
 		}
 	}
-	return &MessageQueueTokenExpiredError{
-		activityID: resp.Header.Get(headerActionsActivityID),
-		statusCode: resp.StatusCode,
-		msg:        string(body),
+	return &ActionsError{
+		ActivityID: resp.Header.Get(headerActionsActivityID),
+		StatusCode: resp.StatusCode,
+		Err: &messageQueueTokenExpiredError{
+			message: string(body),
+		},
 	}
 }
 
@@ -754,13 +758,19 @@ func (c *Client) CreateMessageSession(ctx context.Context, runnerScaleSetID int,
 		return nil, fmt.Errorf("failed to marshal new session: %w", err)
 	}
 
-	createdSession := &RunnerScaleSetSession{}
-
-	if err = c.doSessionRequest(ctx, http.MethodPost, path, bytes.NewBuffer(requestData), http.StatusOK, createdSession); err != nil {
+	var createdSession RunnerScaleSetSession
+	if err = c.doSessionRequest(
+		ctx,
+		http.MethodPost,
+		path,
+		bytes.NewBuffer(requestData),
+		http.StatusOK,
+		&createdSession,
+	); err != nil {
 		return nil, fmt.Errorf("failed to do the session request: %w", err)
 	}
 
-	return createdSession, nil
+	return &createdSession, nil
 }
 
 // DeleteMessageSession deletes a message session for the specified runner scale set.
@@ -809,7 +819,7 @@ func (c *Client) doSessionRequest(ctx context.Context, method, path string, requ
 	}
 
 	if resp.StatusCode >= 400 && resp.StatusCode < 500 {
-		return parseActionsErrorFromResponse(resp)
+		return ParseActionsErrorFromResponse(resp)
 	}
 
 	body, err := io.ReadAll(resp.Body)
@@ -851,7 +861,7 @@ func (c *Client) GenerateJitRunnerConfig(ctx context.Context, jitRunnerSetting *
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var runnerJitConfig *RunnerScaleSetJitRunnerConfig
@@ -881,7 +891,7 @@ func (c *Client) GetRunner(ctx context.Context, runnerID int) (*RunnerReference,
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var runnerReference *RunnerReference
@@ -912,7 +922,7 @@ func (c *Client) GetRunnerByName(ctx context.Context, runnerName string) (*Runne
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, parseActionsErrorFromResponse(resp)
+		return nil, ParseActionsErrorFromResponse(resp)
 	}
 
 	var runnerList *RunnerReferenceList
@@ -955,7 +965,7 @@ func (c *Client) RemoveRunner(ctx context.Context, runnerID int64) error {
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusNoContent {
-		return parseActionsErrorFromResponse(resp)
+		return ParseActionsErrorFromResponse(resp)
 	}
 
 	return nil

client_test.go

@@ -945,7 +945,7 @@ func TestCreateMessageSession(t *testing.T) {
 		want := &ActionsError{
 			ActivityID: exampleRequestID,
 			StatusCode: http.StatusBadRequest,
-			Err: &ActionsExceptionError{
+			Err: &actionsExceptionError{
 				ExceptionName: "CSharpExceptionNameHere",
 				Message:       "could not do something",
 			},
@@ -1169,8 +1169,10 @@ func TestGetMessage(t *testing.T) {
 		_, err = client.GetMessage(ctx, server.URL, token, 0, 10)
 		require.NotNil(t, err)
 
-		var expectedErr *MessageQueueTokenExpiredError
+		var expectedErr *ActionsError
 		require.True(t, errors.As(err, &expectedErr))
+
+		assert.True(t, expectedErr.IsMessageQueueTokenExpired())
 	})
 
 	t.Run("Status code not found", func(t *testing.T) {
@@ -1258,8 +1260,9 @@ func TestDeleteMessage(t *testing.T) {
 
 		err = client.DeleteMessage(ctx, server.URL, token, 0)
 		require.NotNil(t, err)
-		var expectedErr *MessageQueueTokenExpiredError
-		assert.True(t, errors.As(err, &expectedErr))
+		var expectedErr *ActionsError
+		require.ErrorAs(t, err, &expectedErr)
+		assert.True(t, expectedErr.IsMessageQueueTokenExpired())
 	})
 
 	t.Run("Error when Content-Type is text/plain", func(t *testing.T) {

errors.go

@@ -43,23 +43,46 @@ func (e *ActionsError) Unwrap() error {
 	return e.Err
 }
 
-func (e *ActionsError) IsException(target string) bool {
-	if ex, ok := e.Err.(*ActionsExceptionError); ok {
+func (e *ActionsError) IsAgentNotFound() bool {
+	return e.isException("AgentNotFoundException")
+}
+
+func (e *ActionsError) IsJobStillRunning() bool {
+	return e.isException("JobStillRunningException")
+}
+
+func (e *ActionsError) IsMessageQueueTokenExpired() bool {
+	if e == nil {
+		return false
+	}
+	var err *messageQueueTokenExpiredError
+	return errors.As(e.Err, &err)
+}
+
+func (e *ActionsError) IsAgentExists() bool {
+	return e.isException("AgentExistsException")
+}
+
+func (e *ActionsError) isException(target string) bool {
+	if e == nil {
+		return false
+	}
+	if ex, ok := e.Err.(*actionsExceptionError); ok {
 		return strings.Contains(ex.ExceptionName, target)
 	}
 	return false
 }
 
-type ActionsExceptionError struct {
+type actionsExceptionError struct {
 	ExceptionName string `json:"typeName,omitempty"`
 	Message       string `json:"message,omitempty"`
 }
 
-func (e *ActionsExceptionError) Error() string {
+func (e *actionsExceptionError) Error() string {
 	return fmt.Sprintf("%s: %s", e.ExceptionName, e.Message)
 }
 
-func parseActionsErrorFromResponse(response *http.Response) error {
+func ParseActionsErrorFromResponse(response *http.Response) error {
 	if response.ContentLength == 0 {
 		return &ActionsError{
 			ActivityID: response.Header.Get(headerActionsActivityID),
@@ -78,8 +101,8 @@ func parseActionsErrorFromResponse(response *http.Response) error {
 	}
 
 	body = trimByteOrderMark(body)
-	contentType, ok := response.Header["Content-Type"]
-	if ok && len(contentType) > 0 && strings.Contains(contentType[0], "text/plain") {
+	contentType := response.Header.Get("Content-Type")
+	if len(contentType) > 0 && strings.Contains(contentType, "text/plain") {
 		message := string(body)
 		return &ActionsError{
 			ActivityID: response.Header.Get(headerActionsActivityID),
@@ -88,7 +111,7 @@ func parseActionsErrorFromResponse(response *http.Response) error {
 		}
 	}
 
-	var exception ActionsExceptionError
+	var exception actionsExceptionError
 	if err := json.Unmarshal(body, &exception); err != nil {
 		return &ActionsError{
 			ActivityID: response.Header.Get(headerActionsActivityID),
@@ -104,21 +127,19 @@ func parseActionsErrorFromResponse(response *http.Response) error {
 	}
 }
 
-type MessageQueueTokenExpiredError struct {
-	activityID string
-	statusCode int
-	msg        string
+type messageQueueTokenExpiredError struct {
+	message string
 }
 
-func (e *MessageQueueTokenExpiredError) Error() string {
-	return fmt.Sprintf("MessageQueueTokenExpiredError: ActivityId %q, StatusCode %d: %s", e.activityID, e.statusCode, e.msg)
+func (e *messageQueueTokenExpiredError) Error() string {
+	return fmt.Sprintf("message queue token expired: %s", e.message)
 }
 
-type HttpClientSideError struct {
-	msg  string
-	Code int
-}
-
-func (e *HttpClientSideError) Error() string {
-	return e.msg
+// NewMessageQueueTokenExpiredError creates a new MessageQueueTokenExpiredError.
+//
+// This function is mostly used by tests.
+func NewMessageQueueTokenExpiredError(message string) error {
+	return &messageQueueTokenExpiredError{
+		message: message,
+	}
 }

errors_test.go

@@ -29,7 +29,7 @@ func TestActionsError(t *testing.T) {
 		err := &ActionsError{
 			ActivityID: "activity-id",
 			StatusCode: 404,
-			Err: &ActionsExceptionError{
+			Err: &actionsExceptionError{
 				ExceptionName: "exception-name",
 				Message:       "example error message",
 			},
@@ -42,16 +42,16 @@ func TestActionsError(t *testing.T) {
 		err := &ActionsError{
 			ActivityID: "activity-id",
 			StatusCode: 404,
-			Err: &ActionsExceptionError{
+			Err: &actionsExceptionError{
 				ExceptionName: "exception-name",
 				Message:       "example error message",
 			},
 		}
 
-		var exception *ActionsExceptionError
+		var exception *actionsExceptionError
 		assert.True(t, errors.As(err, &exception))
 
-		assert.True(t, err.IsException("exception-name"))
+		assert.True(t, err.isException("exception-name"))
 	})
 
 	t.Run("is exception is not ok", func(t *testing.T) {
@@ -64,7 +64,7 @@ func TestActionsError(t *testing.T) {
 			"not target exception": {
 				ActivityID: "activity-id",
 				StatusCode: 404,
-				Err: &ActionsExceptionError{
+				Err: &actionsExceptionError{
 					ExceptionName: "exception-name",
 					Message:       "example error message",
 				},
@@ -74,15 +74,66 @@ func TestActionsError(t *testing.T) {
 		targetException := "target-exception"
 		for name, err := range tt {
 			t.Run(name, func(t *testing.T) {
-				assert.False(t, err.IsException(targetException))
+				assert.False(t, err.isException(targetException))
 			})
 		}
 	})
+
+	t.Run("is agent exists exception", func(t *testing.T) {
+		err := &ActionsError{
+			ActivityID: "activity-id",
+			StatusCode: 404,
+			Err: &actionsExceptionError{
+				ExceptionName: "AgentExistsException",
+				Message:       "example error message",
+			},
+		}
+
+		assert.True(t, err.IsAgentExists())
+	})
+
+	t.Run("is agent not found exception", func(t *testing.T) {
+		err := &ActionsError{
+			ActivityID: "activity-id",
+			StatusCode: 404,
+			Err: &actionsExceptionError{
+				ExceptionName: "AgentNotFoundException",
+				Message:       "example error message",
+			},
+		}
+
+		assert.True(t, err.IsAgentNotFound())
+	})
+
+	t.Run("is job still running exception", func(t *testing.T) {
+		err := &ActionsError{
+			ActivityID: "activity-id",
+			StatusCode: 404,
+			Err: &actionsExceptionError{
+				ExceptionName: "JobStillRunningException",
+				Message:       "example error message",
+			},
+		}
+
+		assert.True(t, err.IsJobStillRunning())
+	})
+
+	t.Run("is message queue token expired exception", func(t *testing.T) {
+		err := &ActionsError{
+			ActivityID: "activity-id",
+			StatusCode: 404,
+			Err: &messageQueueTokenExpiredError{
+				message: "example error message",
+			},
+		}
+
+		assert.True(t, err.IsMessageQueueTokenExpired())
+	})
 }
 
 func TestActionsExceptionError(t *testing.T) {
 	t.Run("contains the exception name and message", func(t *testing.T) {
-		err := &ActionsExceptionError{
+		err := &actionsExceptionError{
 			ExceptionName: "exception-name",
 			Message:       "example error message",
 		}
@@ -127,7 +178,7 @@ func TestParseActionsErrorFromResponse(t *testing.T) {
 		}
 		response.Header.Add(headerActionsActivityID, "activity-id")
 
-		err := parseActionsErrorFromResponse(response)
+		err := ParseActionsErrorFromResponse(response)
 		require.Error(t, err)
 		assert.Equal(t, "activity-id", err.(*ActionsError).ActivityID)
 		assert.Equal(t, 404, err.(*ActionsError).StatusCode)
@@ -145,7 +196,7 @@ func TestParseActionsErrorFromResponse(t *testing.T) {
 		response.Header.Add(headerActionsActivityID, "activity-id")
 		response.Header.Add("Content-Type", "text/plain")
 
-		err := parseActionsErrorFromResponse(response)
+		err := ParseActionsErrorFromResponse(response)
 		require.Error(t, err)
 		var actionsError *ActionsError
 		assert.ErrorAs(t, err, &actionsError)
@@ -165,14 +216,14 @@ func TestParseActionsErrorFromResponse(t *testing.T) {
 		response.Header.Add(headerActionsActivityID, "activity-id")
 		response.Header.Add("Content-Type", "application/json")
 
-		err := parseActionsErrorFromResponse(response)
+		err := ParseActionsErrorFromResponse(response)
 		require.Error(t, err)
 		var actionsError *ActionsError
 		assert.ErrorAs(t, err, &actionsError)
 		assert.Equal(t, "activity-id", actionsError.ActivityID)
 		assert.Equal(t, 404, actionsError.StatusCode)
 
-		inner, ok := actionsError.Err.(*ActionsExceptionError)
+		inner, ok := actionsError.Err.(*actionsExceptionError)
 		require.True(t, ok)
 		assert.Equal(t, "exception-name", inner.ExceptionName)
 		assert.Equal(t, "example error message", inner.Message)
@@ -189,10 +240,10 @@ func TestParseActionsErrorFromResponse(t *testing.T) {
 		response.Header.Add(headerActionsActivityID, "activity-id")
 		response.Header.Add("Content-Type", "application/json")
 
-		err := parseActionsErrorFromResponse(response)
+		err := ParseActionsErrorFromResponse(response)
 		require.Error(t, err)
 
-		var actionsExceptionError *ActionsExceptionError
+		var actionsExceptionError *actionsExceptionError
 		assert.ErrorAs(t, err, &actionsExceptionError)
 
 		assert.Equal(t, "exception-name", actionsExceptionError.ExceptionName)

examples/dockerscaleset/main.go

@@ -140,7 +140,6 @@ func run(ctx context.Context, c Config) error {
 	logger.Info("Initializing listener")
 	listener, err := listener.New(scalesetClient, listener.Config{
 		ScaleSetID: scaleSet.ID,
-		MinRunners: c.MinRunners,
 		MaxRunners: c.MaxRunners,
 		Logger:     logger.WithGroup("listener"),
 	})

go.mod

@@ -5,6 +5,7 @@ go 1.25.3
 require (
 	github.com/docker/docker v28.5.2+incompatible
 	github.com/golang-jwt/jwt/v4 v4.5.2
+	github.com/google/go-github/v79 v79.0.0
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-retryablehttp v0.7.8
 	github.com/spf13/cobra v1.10.1
@@ -14,6 +15,7 @@ require (
 
 require (
 	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/brunoga/deep v1.2.4 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
@@ -21,13 +23,29 @@ require (
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/go-connections v0.6.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
+	github.com/fatih/structs v1.1.0 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
+	github.com/fsnotify/fsnotify v1.8.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
-	github.com/google/go-github/v79 v79.0.0 // indirect
+	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
 	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
+	github.com/huandu/xstrings v1.5.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jedib0t/go-pretty/v6 v6.6.7 // indirect
+	github.com/knadh/koanf/maps v0.1.2 // indirect
+	github.com/knadh/koanf/parsers/yaml v0.1.0 // indirect
+	github.com/knadh/koanf/providers/env v1.0.0 // indirect
+	github.com/knadh/koanf/providers/file v1.1.2 // indirect
+	github.com/knadh/koanf/providers/posflag v0.1.0 // indirect
+	github.com/knadh/koanf/providers/structs v0.1.0 // indirect
+	github.com/knadh/koanf/v2 v2.3.0 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-runewidth v0.0.16 // indirect
+	github.com/mitchellh/copystructure v1.2.0 // indirect
+	github.com/mitchellh/reflectwalk v1.0.2 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
 	github.com/moby/sys/atomicwriter v0.1.0 // indirect
 	github.com/moby/term v0.5.2 // indirect
@@ -36,17 +54,26 @@ require (
 	github.com/opencontainers/image-spec v1.1.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/rivo/uniseg v0.4.7 // indirect
+	github.com/rs/zerolog v1.33.0 // indirect
 	github.com/spf13/pflag v1.0.10 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
+	github.com/vektra/mockery/v3 v3.6.1 // indirect
+	github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect
+	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
+	github.com/xeipuuv/gojsonschema v1.2.0 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
 	go.opentelemetry.io/otel v1.35.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 // indirect
 	go.opentelemetry.io/otel/metric v1.35.0 // indirect
 	go.opentelemetry.io/otel/trace v1.35.0 // indirect
+	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
 	golang.org/x/mod v0.30.0 // indirect
 	golang.org/x/sync v0.18.0 // indirect
 	golang.org/x/sys v0.38.0 // indirect
 	golang.org/x/telemetry v0.0.0-20251111182119-bc8e575c7b54 // indirect
+	golang.org/x/term v0.37.0 // indirect
 	golang.org/x/text v0.31.0 // indirect
 	golang.org/x/time v0.14.0 // indirect
 	golang.org/x/tools v0.39.0 // indirect
@@ -56,4 +83,7 @@ require (
 	gotest.tools/v3 v3.5.2 // indirect
 )
 
-tool golang.org/x/tools/cmd/deadcode
+tool (
+	github.com/vektra/mockery/v3
+	golang.org/x/tools/cmd/deadcode
+)

go.sum

@@ -2,6 +2,8 @@ github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEK
 github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/brunoga/deep v1.2.4 h1:Aj9E9oUbE+ccbyh35VC/NHlzzjfIVU69BXu2mt2LmL8=
+github.com/brunoga/deep v1.2.4/go.mod h1:GDV6dnXqn80ezsLSZ5Wlv1PdKAWAO4L5PnKYtv2dgaI=
 github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
@@ -10,7 +12,9 @@ github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151X
 github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk=
 github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
 github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
+github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
@@ -23,13 +27,20 @@ github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4
 github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM=
 github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE=
+github.com/fatih/structs v1.1.0 h1:Q7juDM0QtcnhCpeyLGQKyg4TOIghuNXrkL32pHAUMxo=
+github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
 github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
 github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
+github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
 github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
 github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs=
+github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
 github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
@@ -49,16 +60,42 @@ github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB1
 github.com/hashicorp/go-hclog v1.6.3/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M=
 github.com/hashicorp/go-retryablehttp v0.7.8 h1:ylXZWnqa7Lhqpk0L1P1LzDtGcCR0rPVUrx/c8Unxc48=
 github.com/hashicorp/go-retryablehttp v0.7.8/go.mod h1:rjiScheydd+CxvumBsIrFKlx3iS0jrZ7LvzFGFmuKbw=
+github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI=
+github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jedib0t/go-pretty/v6 v6.6.7 h1:m+LbHpm0aIAPLzLbMfn8dc3Ht8MW7lsSO4MPItz/Uuo=
+github.com/jedib0t/go-pretty/v6 v6.6.7/go.mod h1:YwC5CE4fJ1HFUDeivSV1r//AmANFHyqczZk+U6BDALU=
+github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
+github.com/knadh/koanf/maps v0.1.2/go.mod h1:npD/QZY3V6ghQDdcQzl1W4ICNVTkohC8E73eI2xW4yI=
+github.com/knadh/koanf/parsers/yaml v0.1.0 h1:ZZ8/iGfRLvKSaMEECEBPM1HQslrZADk8fP1XFUxVI5w=
+github.com/knadh/koanf/parsers/yaml v0.1.0/go.mod h1:cvbUDC7AL23pImuQP0oRw/hPuccrNBS2bps8asS0CwY=
+github.com/knadh/koanf/providers/env v1.0.0 h1:ufePaI9BnWH+ajuxGGiJ8pdTG0uLEUWC7/HDDPGLah0=
+github.com/knadh/koanf/providers/env v1.0.0/go.mod h1:mzFyRZueYhb37oPmC1HAv/oGEEuyvJDA98r3XAa8Gak=
+github.com/knadh/koanf/providers/file v1.1.2 h1:aCC36YGOgV5lTtAFz2qkgtWdeQsgfxUkxDOe+2nQY3w=
+github.com/knadh/koanf/providers/file v1.1.2/go.mod h1:/faSBcv2mxPVjFrXck95qeoyoZ5myJ6uxN8OOVNJJCI=
+github.com/knadh/koanf/providers/posflag v0.1.0 h1:mKJlLrKPcAP7Ootf4pBZWJ6J+4wHYujwipe7Ie3qW6U=
+github.com/knadh/koanf/providers/posflag v0.1.0/go.mod h1:SYg03v/t8ISBNrMBRMlojH8OsKowbkXV7giIbBVgbz0=
+github.com/knadh/koanf/providers/structs v0.1.0 h1:wJRteCNn1qvLtE5h8KQBvLJovidSdntfdyIbbCzEyE0=
+github.com/knadh/koanf/providers/structs v0.1.0/go.mod h1:sw2YZ3txUcqA3Z27gPlmmBzWn1h8Nt9O6EP/91MkcWE=
+github.com/knadh/koanf/v2 v2.3.0 h1:Qg076dDRFHvqnKG97ZEsi9TAg2/nFTa9hCdcSa1lvlM=
+github.com/knadh/koanf/v2 v2.3.0/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
 github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
+github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
+github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
+github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
+github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
 github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
 github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
 github.com/moby/sys/atomicwriter v0.1.0 h1:kw5D/EqkBwsBFi0ss9v1VG3wIkVhzGvLklJ+w3A14Sw=
@@ -75,10 +112,17 @@ github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJw
 github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
+github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
+github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
+github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
@@ -87,8 +131,20 @@ github.com/spf13/cobra v1.10.1/go.mod h1:7SmJGaTHFVBY0jW4NXGluQoLvhqFQM+6XSKD+P4
 github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
 github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+github.com/vektra/mockery/v3 v3.6.1 h1:YyqAXihdNML8y6SJnvPKYr+2HAHvBjdvqFu/fMYlX8g=
+github.com/vektra/mockery/v3 v3.6.1/go.mod h1:Oti3Df0WP8wwT31yuVri3QNsDeMUQU5Q4QEg8EabaBw=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ=
+github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74=
+github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
 go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
 go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU=
@@ -109,16 +165,23 @@ go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt
 go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc=
 go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU/3i4=
 go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4=
+golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 h1:2dVuKD2vS7b0QIHQbpyTISPd0LeHDbnYEryqj5Q1ug8=
+golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY=
 golang.org/x/mod v0.30.0 h1:fDEXFVZ/fmCKProc/yAXXUijritrDzahmwwefnjoPFk=
 golang.org/x/mod v0.30.0/go.mod h1:lAsf5O2EvJeSFMiBxXDki7sCgAxEUcZHXoXMKT4GJKc=
 golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
 golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
 golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
 golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
 golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
 golang.org/x/telemetry v0.0.0-20251111182119-bc8e575c7b54 h1:E2/AqCUMZGgd73TQkxUMcMla25GB9i/5HOdLr+uH7Vo=
 golang.org/x/telemetry v0.0.0-20251111182119-bc8e575c7b54/go.mod h1:hKdjCMrbv9skySur+Nek8Hd0uJ0GuxJIoIX2payrIdQ=
+golang.org/x/term v0.37.0 h1:8EGAD0qCmHYZg6J17DvsMy9/wJ7/D/4pV/wfnld5lTU=
+golang.org/x/term v0.37.0/go.mod h1:5pB4lxRNYYVZuTLmy8oR2BH8dflOR+IbTYFD8fi3254=
 golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
 golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
 golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=

listener/listener.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"math"
 	"net/http"
 	"os"
 	"sync/atomic"
@@ -22,10 +23,12 @@ const (
 
 // Config holds the configuration for the Listener.
 type Config struct {
+	// ScaleSetID is the ID of the runner scale set to listen to.
 	ScaleSetID int
-	MinRunners int
+	// MaxRunners is the capacity of runners that can be handled at once.
 	MaxRunners int
-	Logger     *slog.Logger
+	// Logger is the logger to use for logging. Default is a no-op logger.
+	Logger *slog.Logger
 }
 
 func (c *Config) defaults() {
@@ -41,17 +44,29 @@ func (c *Config) Validate() error {
 	if c.ScaleSetID == 0 {
 		return errors.New("scaleSetID is required")
 	}
-	if c.MaxRunners > 0 && c.MinRunners > c.MaxRunners {
-		return errors.New("minRunners must be less than or equal to maxRunners")
+	if c.MaxRunners < 0 || c.MaxRunners > math.MaxInt32 {
+		return errors.New("maxRunners must be between 0 and MaxInt32")
 	}
 	return nil
 }
 
+// Client defines the interface for communicating with the scaleset API.
+// In most cases, it should be scaleset.Client from the scaleset package.
+// This interface is defined to allow for easier testing and mocking, as well
+// as allowing wrappers around the scaleset client if needed.
+type Client interface {
+	CreateMessageSession(ctx context.Context, runnerScaleSetID int, owner string) (*scaleset.RunnerScaleSetSession, error)
+	GetMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)
+	DeleteMessage(ctx context.Context, messageQueueURL, messageQueueAccessToken string, messageID int) error
+	RefreshMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) (*scaleset.RunnerScaleSetSession, error)
+	DeleteMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) error
+}
+
 // Listener listens for messages from the scaleset service and handles them. It automatically handles session
 // creation/deletion/refreshing and message polling and acking.
 type Listener struct {
 	// The main client responsible for communicating with the scaleset service
-	client *scaleset.Client
+	client Client
 
 	// Configuration for the listener
 	scaleSetID int
@@ -68,12 +83,14 @@ type Listener struct {
 	logger *slog.Logger
 }
 
+// SetMaxRunners sets the capacity of the scaleset. It is concurrently
+// safe to update the max runners during listener.Run.
 func (l *Listener) SetMaxRunners(count int) {
 	l.maxRunners.Store(uint32(count))
 }
 
 // New creates a new Listener with the given configuration.
-func New(client *scaleset.Client, config Config) (*Listener, error) {
+func New(client Client, config Config) (*Listener, error) {
 	if client == nil {
 		return nil, errors.New("client is required")
 	}
@@ -194,12 +211,12 @@ func (l *Listener) createSession(ctx context.Context) error {
 			break
 		}
 
-		clientErr := &scaleset.HttpClientSideError{}
+		clientErr := &scaleset.ActionsError{}
 		if !errors.As(err, &clientErr) {
 			return fmt.Errorf("failed to create session: %w", err)
 		}
 
-		if clientErr.Code != http.StatusConflict {
+		if clientErr.StatusCode != http.StatusConflict {
 			return fmt.Errorf("failed to create session: %w", err)
 		}
 
@@ -241,8 +258,8 @@ func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMess
 		return msg, nil
 	}
 
-	expiredError := &scaleset.MessageQueueTokenExpiredError{}
-	if !errors.As(err, &expiredError) {
+	expiredError := &scaleset.ActionsError{}
+	if !errors.As(err, &expiredError) || !expiredError.IsMessageQueueTokenExpired() {
 		return nil, fmt.Errorf("failed to get next message: %w", err)
 	}
 
@@ -278,8 +295,8 @@ func (l *Listener) deleteLastMessage(ctx context.Context) error {
 		return nil
 	}
 
-	expiredError := &scaleset.MessageQueueTokenExpiredError{}
-	if !errors.As(err, &expiredError) {
+	expiredError := &scaleset.ActionsError{}
+	if !errors.As(err, &expiredError) || !expiredError.IsMessageQueueTokenExpired() {
 		return fmt.Errorf("failed to delete last message: %w", err)
 	}
 

listener/listener_test.go

@@ -0,0 +1,779 @@
+package listener
+
+import (
+	"context"
+	"errors"
+	"math"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/actions/scaleset"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNew(t *testing.T) {
+	t.Parallel()
+	t.Run("invalid config", func(t *testing.T) {
+		t.Parallel()
+		var config Config
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("valid config", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+		}
+		assert.NoError(t, config.Validate())
+	})
+
+	t.Run("invalid max runners", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: -1,
+		}
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("zero max runners", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: math.MaxInt32 + 1,
+		}
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("creates listener", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 5,
+		}
+
+		client := NewMockClient(t)
+		l, err := New(client, config)
+		require.Nil(t, err)
+		assert.Equal(t, config.ScaleSetID, l.scaleSetID)
+		assert.Equal(t, uint32(config.MaxRunners), l.maxRunners.Load())
+	})
+}
+
+func TestListener_createSession(t *testing.T) {
+	t.Parallel()
+	t.Run("fail once", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+		client.On(
+			"CreateMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(
+			nil,
+			assert.AnError,
+		).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		err = l.createSession(ctx)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("fail context", func(t *testing.T) {
+		t.Parallel()
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+		client.On(
+			"CreateMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(
+			nil,
+			scaleset.ParseActionsErrorFromResponse(&http.Response{
+				StatusCode: http.StatusConflict,
+			}),
+		).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		err = l.createSession(ctx)
+		assert.True(t, errors.Is(err, context.DeadlineExceeded))
+	})
+
+	t.Run("sets session", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		uuid := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"CreateMessageSession",
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		err = l.createSession(context.Background())
+		assert.Nil(t, err)
+		assert.Equal(t, session, l.session)
+	})
+}
+
+func TestListener_getMessage(t *testing.T) {
+	t.Parallel()
+
+	t.Run("receives message", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+		want := &scaleset.RunnerScaleSetMessage{
+			MessageID: 1,
+		}
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).Return(
+			want,
+			nil,
+		).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+		l.session = &scaleset.RunnerScaleSetSession{}
+
+		got, err := l.getMessage(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("not expired error", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).Return(
+			nil,
+			scaleset.ParseActionsErrorFromResponse(&http.Response{
+				StatusCode: http.StatusNotFound,
+			}),
+		).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		l.session = &scaleset.RunnerScaleSetSession{}
+
+		_, err = l.getMessage(ctx)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("refresh and succeeds", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+
+		uuid := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).Return(
+			nil,
+			&scaleset.ActionsError{
+				StatusCode: http.StatusUnauthorized,
+				ActivityID: "1234",
+				Err:        scaleset.NewMessageQueueTokenExpiredError("token expired"),
+			},
+		).Once()
+
+		want := &scaleset.RunnerScaleSetMessage{
+			MessageID: 1,
+		}
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).Return(want, nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		l.session = &scaleset.RunnerScaleSetSession{
+			SessionID:      uuid,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+
+		got, err := l.getMessage(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("refresh and fails", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+
+		uuid := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).Return(
+			nil,
+			&scaleset.ActionsError{
+				StatusCode: http.StatusUnauthorized,
+				ActivityID: "1234",
+				Err:        scaleset.NewMessageQueueTokenExpiredError("token expired"),
+			},
+		).Twice()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		l.session = &scaleset.RunnerScaleSetSession{
+			SessionID:      uuid,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+
+		got, err := l.getMessage(ctx)
+		assert.NotNil(t, err)
+		assert.Nil(t, got)
+	})
+}
+
+func TestListener_refreshSession(t *testing.T) {
+	t.Parallel()
+
+	t.Run("successfully refreshes", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		newUUID := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               newUUID,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		l.session = &scaleset.RunnerScaleSetSession{
+			SessionID:      oldUUID,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+
+		err = l.refreshSession(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, session, l.session)
+	})
+
+	t.Run("fails to refresh", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(nil, errors.New("error")).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		oldSession := &scaleset.RunnerScaleSetSession{
+			SessionID:      oldUUID,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+		l.session = oldSession
+
+		err = l.refreshSession(ctx)
+		assert.NotNil(t, err)
+		assert.Equal(t, oldSession, l.session)
+	})
+}
+
+func TestListener_deleteLastMessage(t *testing.T) {
+	t.Parallel()
+
+	t.Run("successfully deletes", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		client.On(
+			"DeleteMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.MatchedBy(
+				func(lastMessageID any) bool {
+					return lastMessageID.(int) == 5
+				},
+			),
+		).Return(nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		l.session = &scaleset.RunnerScaleSetSession{}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.Nil(t, err)
+	})
+
+	t.Run("fails to delete", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		client.On(
+			"DeleteMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(errors.New("error")).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		l.session = &scaleset.RunnerScaleSetSession{}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("refresh and succeeds", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		newUUID := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               newUUID,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		client.On(
+			"DeleteMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(
+			&scaleset.ActionsError{
+				StatusCode: http.StatusUnauthorized,
+				ActivityID: "1234",
+				Err:        scaleset.NewMessageQueueTokenExpiredError("token expired"),
+			},
+		).Once()
+
+		client.On(
+			"DeleteMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.MatchedBy(
+				func(lastMessageID any) bool {
+					return lastMessageID.(int) == 5
+				},
+			),
+		).Return(nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		l.session = &scaleset.RunnerScaleSetSession{
+			SessionID:      oldUUID,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.NoError(t, err)
+	})
+
+	t.Run("refresh and fails", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		newUUID := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               newUUID,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              nil,
+		}
+		client.On(
+			"RefreshMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+
+		client.On(
+			"DeleteMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(
+			&scaleset.ActionsError{
+				StatusCode: http.StatusUnauthorized,
+				ActivityID: "1234",
+				Err:        scaleset.NewMessageQueueTokenExpiredError("token expired"),
+			},
+		).Twice()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		oldUUID := uuid.New()
+		l.session = &scaleset.RunnerScaleSetSession{
+			SessionID:      oldUUID,
+			RunnerScaleSet: &scaleset.RunnerScaleSet{},
+		}
+		l.lastMessageID = 5
+
+		err = l.deleteLastMessage(ctx)
+		assert.Error(t, err)
+	})
+}
+
+func TestListener_Run(t *testing.T) {
+	t.Parallel()
+
+	t.Run("create session fails", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+		client.On(
+			"CreateMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(nil, assert.AnError).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		err = l.Run(ctx, nil)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("call handle regardless of initial message", func(t *testing.T) {
+		t.Parallel()
+		ctx, cancel := context.WithCancel(context.Background())
+
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		uuid := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              &scaleset.RunnerScaleSetStatistic{},
+		}
+		client.On(
+			"CreateMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+		client.On(
+			"DeleteMessageSession",
+			mock.Anything,
+			session.RunnerScaleSet.ID,
+			session.SessionID,
+		).Return(nil).Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		var called bool
+		handler := NewMockScaler(t)
+		handler.On(
+			"HandleDesiredRunnerCount",
+			mock.Anything,
+			mock.Anything,
+		).
+			Return(0, nil).
+			Run(
+				func(mock.Arguments) {
+					called = true
+					cancel()
+				},
+			).
+			Once()
+
+		err = l.Run(ctx, handler)
+		assert.ErrorIs(t, err, context.Canceled)
+		assert.True(t, called)
+	})
+
+	t.Run("cancel context after get message", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(context.Background())
+
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		client := NewMockClient(t)
+		uuid := uuid.New()
+		session := &scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              &scaleset.RunnerScaleSetStatistic{},
+		}
+		client.On(
+			"CreateMessageSession",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+		).Return(session, nil).Once()
+		client.On(
+			"DeleteMessageSession",
+			mock.Anything,
+			session.RunnerScaleSet.ID,
+			session.SessionID,
+		).Return(nil).Once()
+
+		msg := &scaleset.RunnerScaleSetMessage{
+			MessageID:  1,
+			Statistics: &scaleset.RunnerScaleSetStatistic{},
+		}
+		client.On(
+			"GetMessage",
+			ctx,
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+			10,
+		).
+			Return(msg, nil).
+			Run(
+				func(mock.Arguments) {
+					cancel()
+				},
+			).
+			Once()
+
+		// Ensure delete message is called without cancel
+		client.On(
+			"DeleteMessage",
+			context.WithoutCancel(ctx),
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(nil).Once()
+
+		handler := NewMockScaler(t)
+		handler.On(
+			"HandleDesiredRunnerCount",
+			mock.Anything,
+			0,
+		).
+			Return(0, nil).
+			Once()
+
+		handler.On(
+			"HandleDesiredRunnerCount",
+			mock.Anything,
+			mock.Anything,
+		).
+			Return(0, nil).
+			Once()
+
+		l, err := New(client, config)
+		require.Nil(t, err)
+
+		err = l.Run(ctx, handler)
+		assert.ErrorIs(t, context.Canceled, err)
+	})
+}

listener/mocks_test.go

@@ -0,0 +1,613 @@
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package listener
+
+import (
+	"context"
+
+	"github.com/actions/scaleset"
+	"github.com/google/uuid"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// NewMockClient creates a new instance of MockClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockClient {
+	mock := &MockClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// MockClient is an autogenerated mock type for the Client type
+type MockClient struct {
+	mock.Mock
+}
+
+type MockClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockClient) EXPECT() *MockClient_Expecter {
+	return &MockClient_Expecter{mock: &_m.Mock}
+}
+
+// CreateMessageSession provides a mock function for the type MockClient
+func (_mock *MockClient) CreateMessageSession(ctx context.Context, runnerScaleSetID int, owner string) (*scaleset.RunnerScaleSetSession, error) {
+	ret := _mock.Called(ctx, runnerScaleSetID, owner)
+
+	if len(ret) == 0 {
+		panic("no return value specified for CreateMessageSession")
+	}
+
+	var r0 *scaleset.RunnerScaleSetSession
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, string) (*scaleset.RunnerScaleSetSession, error)); ok {
+		return returnFunc(ctx, runnerScaleSetID, owner)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, string) *scaleset.RunnerScaleSetSession); ok {
+		r0 = returnFunc(ctx, runnerScaleSetID, owner)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*scaleset.RunnerScaleSetSession)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, int, string) error); ok {
+		r1 = returnFunc(ctx, runnerScaleSetID, owner)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockClient_CreateMessageSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateMessageSession'
+type MockClient_CreateMessageSession_Call struct {
+	*mock.Call
+}
+
+// CreateMessageSession is a helper method to define mock.On call
+//   - ctx context.Context
+//   - runnerScaleSetID int
+//   - owner string
+func (_e *MockClient_Expecter) CreateMessageSession(ctx interface{}, runnerScaleSetID interface{}, owner interface{}) *MockClient_CreateMessageSession_Call {
+	return &MockClient_CreateMessageSession_Call{Call: _e.mock.On("CreateMessageSession", ctx, runnerScaleSetID, owner)}
+}
+
+func (_c *MockClient_CreateMessageSession_Call) Run(run func(ctx context.Context, runnerScaleSetID int, owner string)) *MockClient_CreateMessageSession_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		var arg2 string
+		if args[2] != nil {
+			arg2 = args[2].(string)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_CreateMessageSession_Call) Return(runnerScaleSetSession *scaleset.RunnerScaleSetSession, err error) *MockClient_CreateMessageSession_Call {
+	_c.Call.Return(runnerScaleSetSession, err)
+	return _c
+}
+
+func (_c *MockClient_CreateMessageSession_Call) RunAndReturn(run func(ctx context.Context, runnerScaleSetID int, owner string) (*scaleset.RunnerScaleSetSession, error)) *MockClient_CreateMessageSession_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// DeleteMessage provides a mock function for the type MockClient
+func (_mock *MockClient) DeleteMessage(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, messageID int) error {
+	ret := _mock.Called(ctx, messageQueueURL, messageQueueAccessToken, messageID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for DeleteMessage")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string, string, int) error); ok {
+		r0 = returnFunc(ctx, messageQueueURL, messageQueueAccessToken, messageID)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockClient_DeleteMessage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteMessage'
+type MockClient_DeleteMessage_Call struct {
+	*mock.Call
+}
+
+// DeleteMessage is a helper method to define mock.On call
+//   - ctx context.Context
+//   - messageQueueURL string
+//   - messageQueueAccessToken string
+//   - messageID int
+func (_e *MockClient_Expecter) DeleteMessage(ctx interface{}, messageQueueURL interface{}, messageQueueAccessToken interface{}, messageID interface{}) *MockClient_DeleteMessage_Call {
+	return &MockClient_DeleteMessage_Call{Call: _e.mock.On("DeleteMessage", ctx, messageQueueURL, messageQueueAccessToken, messageID)}
+}
+
+func (_c *MockClient_DeleteMessage_Call) Run(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, messageID int)) *MockClient_DeleteMessage_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		var arg2 string
+		if args[2] != nil {
+			arg2 = args[2].(string)
+		}
+		var arg3 int
+		if args[3] != nil {
+			arg3 = args[3].(int)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_DeleteMessage_Call) Return(err error) *MockClient_DeleteMessage_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockClient_DeleteMessage_Call) RunAndReturn(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, messageID int) error) *MockClient_DeleteMessage_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// DeleteMessageSession provides a mock function for the type MockClient
+func (_mock *MockClient) DeleteMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) error {
+	ret := _mock.Called(ctx, runnerScaleSetID, sessionID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for DeleteMessageSession")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, uuid.UUID) error); ok {
+		r0 = returnFunc(ctx, runnerScaleSetID, sessionID)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockClient_DeleteMessageSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteMessageSession'
+type MockClient_DeleteMessageSession_Call struct {
+	*mock.Call
+}
+
+// DeleteMessageSession is a helper method to define mock.On call
+//   - ctx context.Context
+//   - runnerScaleSetID int
+//   - sessionID uuid.UUID
+func (_e *MockClient_Expecter) DeleteMessageSession(ctx interface{}, runnerScaleSetID interface{}, sessionID interface{}) *MockClient_DeleteMessageSession_Call {
+	return &MockClient_DeleteMessageSession_Call{Call: _e.mock.On("DeleteMessageSession", ctx, runnerScaleSetID, sessionID)}
+}
+
+func (_c *MockClient_DeleteMessageSession_Call) Run(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID)) *MockClient_DeleteMessageSession_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		var arg2 uuid.UUID
+		if args[2] != nil {
+			arg2 = args[2].(uuid.UUID)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_DeleteMessageSession_Call) Return(err error) *MockClient_DeleteMessageSession_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockClient_DeleteMessageSession_Call) RunAndReturn(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) error) *MockClient_DeleteMessageSession_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetMessage provides a mock function for the type MockClient
+func (_mock *MockClient) GetMessage(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error) {
+	ret := _mock.Called(ctx, messageQueueURL, messageQueueAccessToken, lastMessageID, maxCapacity)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetMessage")
+	}
+
+	var r0 *scaleset.RunnerScaleSetMessage
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string, string, int, int) (*scaleset.RunnerScaleSetMessage, error)); ok {
+		return returnFunc(ctx, messageQueueURL, messageQueueAccessToken, lastMessageID, maxCapacity)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string, string, int, int) *scaleset.RunnerScaleSetMessage); ok {
+		r0 = returnFunc(ctx, messageQueueURL, messageQueueAccessToken, lastMessageID, maxCapacity)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*scaleset.RunnerScaleSetMessage)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, string, string, int, int) error); ok {
+		r1 = returnFunc(ctx, messageQueueURL, messageQueueAccessToken, lastMessageID, maxCapacity)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockClient_GetMessage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetMessage'
+type MockClient_GetMessage_Call struct {
+	*mock.Call
+}
+
+// GetMessage is a helper method to define mock.On call
+//   - ctx context.Context
+//   - messageQueueURL string
+//   - messageQueueAccessToken string
+//   - lastMessageID int
+//   - maxCapacity int
+func (_e *MockClient_Expecter) GetMessage(ctx interface{}, messageQueueURL interface{}, messageQueueAccessToken interface{}, lastMessageID interface{}, maxCapacity interface{}) *MockClient_GetMessage_Call {
+	return &MockClient_GetMessage_Call{Call: _e.mock.On("GetMessage", ctx, messageQueueURL, messageQueueAccessToken, lastMessageID, maxCapacity)}
+}
+
+func (_c *MockClient_GetMessage_Call) Run(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, lastMessageID int, maxCapacity int)) *MockClient_GetMessage_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		var arg2 string
+		if args[2] != nil {
+			arg2 = args[2].(string)
+		}
+		var arg3 int
+		if args[3] != nil {
+			arg3 = args[3].(int)
+		}
+		var arg4 int
+		if args[4] != nil {
+			arg4 = args[4].(int)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+			arg4,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_GetMessage_Call) Return(runnerScaleSetMessage *scaleset.RunnerScaleSetMessage, err error) *MockClient_GetMessage_Call {
+	_c.Call.Return(runnerScaleSetMessage, err)
+	return _c
+}
+
+func (_c *MockClient_GetMessage_Call) RunAndReturn(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)) *MockClient_GetMessage_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RefreshMessageSession provides a mock function for the type MockClient
+func (_mock *MockClient) RefreshMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) (*scaleset.RunnerScaleSetSession, error) {
+	ret := _mock.Called(ctx, runnerScaleSetID, sessionID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RefreshMessageSession")
+	}
+
+	var r0 *scaleset.RunnerScaleSetSession
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, uuid.UUID) (*scaleset.RunnerScaleSetSession, error)); ok {
+		return returnFunc(ctx, runnerScaleSetID, sessionID)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, uuid.UUID) *scaleset.RunnerScaleSetSession); ok {
+		r0 = returnFunc(ctx, runnerScaleSetID, sessionID)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*scaleset.RunnerScaleSetSession)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, int, uuid.UUID) error); ok {
+		r1 = returnFunc(ctx, runnerScaleSetID, sessionID)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockClient_RefreshMessageSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RefreshMessageSession'
+type MockClient_RefreshMessageSession_Call struct {
+	*mock.Call
+}
+
+// RefreshMessageSession is a helper method to define mock.On call
+//   - ctx context.Context
+//   - runnerScaleSetID int
+//   - sessionID uuid.UUID
+func (_e *MockClient_Expecter) RefreshMessageSession(ctx interface{}, runnerScaleSetID interface{}, sessionID interface{}) *MockClient_RefreshMessageSession_Call {
+	return &MockClient_RefreshMessageSession_Call{Call: _e.mock.On("RefreshMessageSession", ctx, runnerScaleSetID, sessionID)}
+}
+
+func (_c *MockClient_RefreshMessageSession_Call) Run(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID)) *MockClient_RefreshMessageSession_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		var arg2 uuid.UUID
+		if args[2] != nil {
+			arg2 = args[2].(uuid.UUID)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_RefreshMessageSession_Call) Return(runnerScaleSetSession *scaleset.RunnerScaleSetSession, err error) *MockClient_RefreshMessageSession_Call {
+	_c.Call.Return(runnerScaleSetSession, err)
+	return _c
+}
+
+func (_c *MockClient_RefreshMessageSession_Call) RunAndReturn(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) (*scaleset.RunnerScaleSetSession, error)) *MockClient_RefreshMessageSession_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockScaler creates a new instance of MockScaler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockScaler(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockScaler {
+	mock := &MockScaler{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// MockScaler is an autogenerated mock type for the Scaler type
+type MockScaler struct {
+	mock.Mock
+}
+
+type MockScaler_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockScaler) EXPECT() *MockScaler_Expecter {
+	return &MockScaler_Expecter{mock: &_m.Mock}
+}
+
+// HandleDesiredRunnerCount provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) {
+	ret := _mock.Called(ctx, count)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleDesiredRunnerCount")
+	}
+
+	var r0 int
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int) (int, error)); ok {
+		return returnFunc(ctx, count)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int) int); ok {
+		r0 = returnFunc(ctx, count)
+	} else {
+		r0 = ret.Get(0).(int)
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, int) error); ok {
+		r1 = returnFunc(ctx, count)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockScaler_HandleDesiredRunnerCount_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleDesiredRunnerCount'
+type MockScaler_HandleDesiredRunnerCount_Call struct {
+	*mock.Call
+}
+
+// HandleDesiredRunnerCount is a helper method to define mock.On call
+//   - ctx context.Context
+//   - count int
+func (_e *MockScaler_Expecter) HandleDesiredRunnerCount(ctx interface{}, count interface{}) *MockScaler_HandleDesiredRunnerCount_Call {
+	return &MockScaler_HandleDesiredRunnerCount_Call{Call: _e.mock.On("HandleDesiredRunnerCount", ctx, count)}
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) Run(run func(ctx context.Context, count int)) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) Return(n int, err error) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Return(n, err)
+	return _c
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) RunAndReturn(run func(ctx context.Context, count int) (int, error)) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// HandleJobCompleted provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error {
+	ret := _mock.Called(ctx, jobInfo)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleJobCompleted")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *scaleset.JobCompleted) error); ok {
+		r0 = returnFunc(ctx, jobInfo)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockScaler_HandleJobCompleted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleJobCompleted'
+type MockScaler_HandleJobCompleted_Call struct {
+	*mock.Call
+}
+
+// HandleJobCompleted is a helper method to define mock.On call
+//   - ctx context.Context
+//   - jobInfo *scaleset.JobCompleted
+func (_e *MockScaler_Expecter) HandleJobCompleted(ctx interface{}, jobInfo interface{}) *MockScaler_HandleJobCompleted_Call {
+	return &MockScaler_HandleJobCompleted_Call{Call: _e.mock.On("HandleJobCompleted", ctx, jobInfo)}
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) Run(run func(ctx context.Context, jobInfo *scaleset.JobCompleted)) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *scaleset.JobCompleted
+		if args[1] != nil {
+			arg1 = args[1].(*scaleset.JobCompleted)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) Return(err error) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) RunAndReturn(run func(ctx context.Context, jobInfo *scaleset.JobCompleted) error) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// HandleJobStarted provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error {
+	ret := _mock.Called(ctx, jobInfo)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleJobStarted")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *scaleset.JobStarted) error); ok {
+		r0 = returnFunc(ctx, jobInfo)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockScaler_HandleJobStarted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleJobStarted'
+type MockScaler_HandleJobStarted_Call struct {
+	*mock.Call
+}
+
+// HandleJobStarted is a helper method to define mock.On call
+//   - ctx context.Context
+//   - jobInfo *scaleset.JobStarted
+func (_e *MockScaler_Expecter) HandleJobStarted(ctx interface{}, jobInfo interface{}) *MockScaler_HandleJobStarted_Call {
+	return &MockScaler_HandleJobStarted_Call{Call: _e.mock.On("HandleJobStarted", ctx, jobInfo)}
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) Run(run func(ctx context.Context, jobInfo *scaleset.JobStarted)) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *scaleset.JobStarted
+		if args[1] != nil {
+			arg1 = args[1].(*scaleset.JobStarted)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) Return(err error) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) RunAndReturn(run func(ctx context.Context, jobInfo *scaleset.JobStarted) error) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Return(run)
+	return _c
+}