actions/scaleset
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Hide session handling and generate session client from the main client (#42)

Browse Source
This commit is contained in:
Nikola Jokic
2026-01-05 18:19:41 +01:00
committed by GitHub
parent ebe67dba45
commit 1208e1216e
11 changed files with 1563 additions and 2092 deletions
Download Patch File Download Diff File Expand all files Collapse all files
listener/listener.go
+30 -196
View File
@@ -3,24 +3,16 @@ package listener
import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"math"
"net/http"
"os"
"sync/atomic"
"time"
"github.com/actions/scaleset"
"github.com/google/uuid"
)
const (
sessionCreationMaxRetries = 10
)
// Config holds the configuration for the Listener.
type Config struct {
// ScaleSetID is the ID of the runner scale set to listen to.
@@ -55,13 +47,13 @@ func (c *Config) Validate() error {
// 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
GetMessage(ctx context.Context, lastMessageID, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)
DeleteMessage(ctx context.Context, messageID int) error
Session() scaleset.RunnerScaleSetSession
}
type Option func(*Listener)
// 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 {
@@ -72,13 +64,6 @@ type Listener struct {
scaleSetID int
maxRunners atomic.Uint32
// lastMessageID keeps track of the last processed message ID
lastMessageID int
// hostname of the current machine
hostname string
// session represents the current message session
session *scaleset.RunnerScaleSetSession
// configuration for the listener
logger *slog.Logger
}
@@ -90,7 +75,7 @@ func (l *Listener) SetMaxRunners(count int) {
}
// New creates a new Listener with the given configuration.
func New(client Client, config Config) (*Listener, error) {
func New(client Client, config Config, options ...Option) (*Listener, error) {
if client == nil {
return nil, errors.New("client is required")
}
@@ -99,16 +84,9 @@ func New(client Client, config Config) (*Listener, error) {
return nil, fmt.Errorf("invalid config: %w", err)
}
hostname, err := os.Hostname()
if err != nil {
hostname = uuid.NewString()
config.Logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
}
listener := &Listener{
client: client,
scaleSetID: config.ScaleSetID,
hostname: hostname,
logger: config.Logger,
}
listener.SetMaxRunners(config.MaxRunners)
@@ -125,32 +103,27 @@ type Scaler interface {
// Run starts the listener and processes messages using the provided scaler.
func (l *Listener) Run(ctx context.Context, scaler Scaler) error {
l.logger.Info("Creating message session")
if err := l.createSession(ctx); err != nil {
return fmt.Errorf("failed to create session: %w", err)
}
{
initialSession := l.client.Session()
defer func() {
l.logger.Debug("Deleting message session")
if err := l.deleteMessageSession(); err != nil {
l.logger.Error(
"failed to delete message session",
slog.String("error", err.Error()),
)
if initialSession.SessionID == uuid.Nil {
return fmt.Errorf("initial session is nil")
}
}()
if l.session.Statistics == nil {
return fmt.Errorf("session statistics is nil")
}
if initialSession.Statistics == nil {
return fmt.Errorf("session statistics is nil")
}
l.logger.Info("Message session created; listening for messages", "sessionID", l.session.SessionID)
// Handle initial statistics
if _, err := scaler.HandleDesiredRunnerCount(ctx, l.session.Statistics.TotalAssignedJobs); err != nil {
return fmt.Errorf("handling initial message failed: %w", err)
l.logger.Info(
"Handling initial session statistics",
slog.Int("totalAssignedJobs", initialSession.Statistics.TotalAssignedJobs),
)
if _, err := scaler.HandleDesiredRunnerCount(ctx, initialSession.Statistics.TotalAssignedJobs); err != nil {
return fmt.Errorf("handling initial message failed: %w", err)
}
}
var lastMessageID int
for {
select {
case <-ctx.Done():
@@ -158,7 +131,12 @@ func (l *Listener) Run(ctx context.Context, scaler Scaler) error {
default:
}
msg, err := l.getMessage(ctx)
l.logger.Info("Getting next message", slog.Int("lastMessageID", lastMessageID))
msg, err := l.client.GetMessage(
ctx,
lastMessageID,
int(l.maxRunners.Load()),
)
if err != nil {
return fmt.Errorf("failed to get message: %w", err)
}
@@ -172,6 +150,8 @@ func (l *Listener) Run(ctx context.Context, scaler Scaler) error {
continue
}
lastMessageID = msg.MessageID
// Remove cancellation from the context to avoid cancelling the message handling.
if err := l.handleMessage(context.WithoutCancel(ctx), scaler, msg); err != nil {
return fmt.Errorf("failed to handle message: %w", err)
@@ -180,8 +160,7 @@ func (l *Listener) Run(ctx context.Context, scaler Scaler) error {
}
func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scaleset.RunnerScaleSetMessage) error {
l.lastMessageID = msg.MessageID
if err := l.deleteLastMessage(ctx); err != nil {
if err := l.client.DeleteMessage(ctx, msg.MessageID); err != nil {
return fmt.Errorf("failed to delete message: %w", err)
}
@@ -202,148 +181,3 @@ func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scale
return nil
}
func (l *Listener) createSession(ctx context.Context) error {
var session *scaleset.RunnerScaleSetSession
var retries int
for {
var err error
session, err = l.client.CreateMessageSession(ctx, l.scaleSetID, l.hostname)
if err == nil {
break
}
clientErr := &scaleset.ActionsError{}
if !errors.As(err, &clientErr) {
return fmt.Errorf("failed to create session: %w", err)
}
if clientErr.StatusCode != http.StatusConflict {
return fmt.Errorf("failed to create session: %w", err)
}
retries++
if retries >= sessionCreationMaxRetries {
return fmt.Errorf("failed to create session after %d retries: %w", retries, err)
}
l.logger.Info("Unable to create message session. Will try again in 30 seconds", "error", err.Error())
select {
case <-ctx.Done():
return fmt.Errorf("context cancelled: %w", ctx.Err())
case <-time.After(30 * time.Second):
}
}
statistics, err := json.Marshal(session.Statistics)
if err != nil {
return fmt.Errorf("failed to marshal statistics: %w", err)
}
l.logger.Info("Current runner scale set statistics.", "statistics", string(statistics))
l.session = session
return nil
}
func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMessage, error) {
l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
msg, err := l.client.GetMessage(
ctx,
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
int(l.maxRunners.Load()),
)
if err == nil { // if NO error
return msg, nil
}
expiredError := &scaleset.ActionsError{}
if !errors.As(err, &expiredError) || !expiredError.IsMessageQueueTokenExpired() {
return nil, fmt.Errorf("failed to get next message: %w", err)
}
if err := l.refreshSession(ctx); err != nil {
return nil, fmt.Errorf("failed to refresh message session: %w", err)
}
l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
msg, err = l.client.GetMessage(
ctx,
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
int(l.maxRunners.Load()),
)
if err != nil { // if error
return nil, fmt.Errorf("failed to get next message after message session refresh: %w", err)
}
return msg, nil
}
func (l *Listener) deleteLastMessage(ctx context.Context) error {
l.logger.Info("Deleting last message", "lastMessageID", l.lastMessageID)
err := l.client.DeleteMessage(
ctx,
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
)
if err == nil { // if NO error
return nil
}
expiredError := &scaleset.ActionsError{}
if !errors.As(err, &expiredError) || !expiredError.IsMessageQueueTokenExpired() {
return fmt.Errorf("failed to delete last message: %w", err)
}
if err := l.refreshSession(ctx); err != nil {
return fmt.Errorf("failed to refresh message session: %w", err)
}
err = l.client.DeleteMessage(
ctx,
l.session.MessageQueueURL,
l.session.MessageQueueAccessToken,
l.lastMessageID,
)
if err != nil {
return fmt.Errorf("failed to delete last message after message session refresh: %w", err)
}
return nil
}
func (l *Listener) refreshSession(ctx context.Context) error {
l.logger.Info("Message queue token is expired during GetNextMessage, refreshing...")
session, err := l.client.RefreshMessageSession(
ctx,
l.session.RunnerScaleSet.ID,
l.session.SessionID,
)
if err != nil {
return fmt.Errorf("refresh message session failed. %w", err)
}
l.session = session
return nil
}
func (l *Listener) deleteMessageSession() error {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
l.logger.Info("Deleting message session")
if err := l.client.DeleteMessageSession(ctx, l.session.RunnerScaleSet.ID, l.session.SessionID); err != nil {
return fmt.Errorf("failed to delete message session: %w", err)
}
return nil
}
listener/listener_test.go
+5 -601
View File
@@ -2,11 +2,8 @@ package listener
import (
"context"
"errors"
"math"
"net/http"
"testing"
"time"
"github.com/actions/scaleset"
"github.com/google/uuid"
@@ -64,577 +61,9 @@ func TestNew(t *testing.T) {
})
}
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())
@@ -646,7 +75,7 @@ func TestListener_Run(t *testing.T) {
client := NewMockClient(t)
uuid := uuid.New()
session := &scaleset.RunnerScaleSetSession{
session := scaleset.RunnerScaleSetSession{
SessionID: uuid,
OwnerName: "example",
RunnerScaleSet: &scaleset.RunnerScaleSet{},
@@ -654,18 +83,8 @@ func TestListener_Run(t *testing.T) {
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()
client.On("Session").Return(session).Once()
l, err := New(client, config)
require.Nil(t, err)
@@ -703,7 +122,7 @@ func TestListener_Run(t *testing.T) {
client := NewMockClient(t)
uuid := uuid.New()
session := &scaleset.RunnerScaleSetSession{
session := scaleset.RunnerScaleSetSession{
SessionID: uuid,
OwnerName: "example",
RunnerScaleSet: &scaleset.RunnerScaleSet{},
@@ -711,29 +130,16 @@ func TestListener_Run(t *testing.T) {
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("Session").Return(session).Once()
client.On(
"GetMessage",
ctx,
mock.Anything,
mock.Anything,
mock.Anything,
10,
).
Return(msg, nil).
@@ -749,8 +155,6 @@ func TestListener_Run(t *testing.T) {
"DeleteMessage",
context.WithoutCancel(ctx),
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil).Once()
handler := NewMockScaler(t)
listener/mocks_test.go
+45 -237
View File
@@ -8,7 +8,6 @@ import (
"context"
"github.com/actions/scaleset"
"github.com/google/uuid"
mock "github.com/stretchr/testify/mock"
)
@@ -39,91 +38,17 @@ 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)
func (_mock *MockClient) DeleteMessage(ctx context.Context, messageID int) error {
ret := _mock.Called(ctx, 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)
if returnFunc, ok := ret.Get(0).(func(context.Context, int) error); ok {
r0 = returnFunc(ctx, messageID)
} else {
r0 = ret.Error(0)
}
@@ -137,36 +62,24 @@ type MockClient_DeleteMessage_Call struct {
// 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 (_e *MockClient_Expecter) DeleteMessage(ctx interface{}, messageID interface{}) *MockClient_DeleteMessage_Call {
return &MockClient_DeleteMessage_Call{Call: _e.mock.On("DeleteMessage", ctx, messageID)}
}
func (_c *MockClient_DeleteMessage_Call) Run(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, messageID int)) *MockClient_DeleteMessage_Call {
func (_c *MockClient_DeleteMessage_Call) Run(run func(ctx context.Context, 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
var arg1 int
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)
arg1 = args[1].(int)
}
run(
arg0,
arg1,
arg2,
arg3,
)
})
return _c
@@ -177,77 +90,14 @@ func (_c *MockClient_DeleteMessage_Call) Return(err error) *MockClient_DeleteMes
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 {
func (_c *MockClient_DeleteMessage_Call) RunAndReturn(run func(ctx context.Context, messageID int) error) *MockClient_DeleteMessage_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)
func (_mock *MockClient) GetMessage(ctx context.Context, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error) {
ret := _mock.Called(ctx, lastMessageID, maxCapacity)
if len(ret) == 0 {
panic("no return value specified for GetMessage")
@@ -255,18 +105,18 @@ func (_mock *MockClient) GetMessage(ctx context.Context, messageQueueURL string,
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, int, int) (*scaleset.RunnerScaleSetMessage, error)); ok {
return returnFunc(ctx, 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)
if returnFunc, ok := ret.Get(0).(func(context.Context, int, int) *scaleset.RunnerScaleSetMessage); ok {
r0 = returnFunc(ctx, 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)
if returnFunc, ok := ret.Get(1).(func(context.Context, int, int) error); ok {
r1 = returnFunc(ctx, lastMessageID, maxCapacity)
} else {
r1 = ret.Error(1)
}
@@ -280,42 +130,30 @@ type MockClient_GetMessage_Call struct {
// 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 (_e *MockClient_Expecter) GetMessage(ctx interface{}, lastMessageID interface{}, maxCapacity interface{}) *MockClient_GetMessage_Call {
return &MockClient_GetMessage_Call{Call: _e.mock.On("GetMessage", ctx, lastMessageID, maxCapacity)}
}
func (_c *MockClient_GetMessage_Call) Run(run func(ctx context.Context, messageQueueURL string, messageQueueAccessToken string, lastMessageID int, maxCapacity int)) *MockClient_GetMessage_Call {
func (_c *MockClient_GetMessage_Call) Run(run func(ctx context.Context, 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
var arg1 int
if args[1] != nil {
arg1 = args[1].(string)
arg1 = args[1].(int)
}
var arg2 string
var arg2 int
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)
arg2 = args[2].(int)
}
run(
arg0,
arg1,
arg2,
arg3,
arg4,
)
})
return _c
@@ -326,81 +164,51 @@ func (_c *MockClient_GetMessage_Call) Return(runnerScaleSetMessage *scaleset.Run
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 {
func (_c *MockClient_GetMessage_Call) RunAndReturn(run func(ctx context.Context, 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)
// Session provides a mock function for the type MockClient
func (_mock *MockClient) Session() scaleset.RunnerScaleSetSession {
ret := _mock.Called()
if len(ret) == 0 {
panic("no return value specified for RefreshMessageSession")
panic("no return value specified for Session")
}
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)
var r0 scaleset.RunnerScaleSetSession
if returnFunc, ok := ret.Get(0).(func() scaleset.RunnerScaleSetSession); ok {
r0 = returnFunc()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*scaleset.RunnerScaleSetSession)
}
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
return r0
}
// MockClient_RefreshMessageSession_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RefreshMessageSession'
type MockClient_RefreshMessageSession_Call struct {
// MockClient_Session_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Session'
type MockClient_Session_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)}
// Session is a helper method to define mock.On call
func (_e *MockClient_Expecter) Session() *MockClient_Session_Call {
return &MockClient_Session_Call{Call: _e.mock.On("Session")}
}
func (_c *MockClient_RefreshMessageSession_Call) Run(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID)) *MockClient_RefreshMessageSession_Call {
func (_c *MockClient_Session_Call) Run(run func()) *MockClient_Session_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,
)
run()
})
return _c
}
func (_c *MockClient_RefreshMessageSession_Call) Return(runnerScaleSetSession *scaleset.RunnerScaleSetSession, err error) *MockClient_RefreshMessageSession_Call {
_c.Call.Return(runnerScaleSetSession, err)
func (_c *MockClient_Session_Call) Return(runnerScaleSetSession scaleset.RunnerScaleSetSession) *MockClient_Session_Call {
_c.Call.Return(runnerScaleSetSession)
return _c
}
func (_c *MockClient_RefreshMessageSession_Call) RunAndReturn(run func(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) (*scaleset.RunnerScaleSetSession, error)) *MockClient_RefreshMessageSession_Call {
func (_c *MockClient_Session_Call) RunAndReturn(run func() scaleset.RunnerScaleSetSession) *MockClient_Session_Call {
_c.Call.Return(run)
return _c
}