forgejo-runner/internal/app/cmd/daemon_test.go

// Copyright 2025 The Forgejo Authors
// SPDX-License-Identifier: MIT

package cmd

import (
	"context"
	"testing"
	"time"

	"code.forgejo.org/forgejo/runner/v11/internal/app/poll"
	mock_poller "code.forgejo.org/forgejo/runner/v11/internal/app/poll/mocks"
	"code.forgejo.org/forgejo/runner/v11/internal/app/run"
	mock_runner "code.forgejo.org/forgejo/runner/v11/internal/app/run/mocks"
	"code.forgejo.org/forgejo/runner/v11/internal/pkg/client"
	mock_client "code.forgejo.org/forgejo/runner/v11/internal/pkg/client/mocks"
	"code.forgejo.org/forgejo/runner/v11/internal/pkg/config"
	"code.forgejo.org/forgejo/runner/v11/internal/pkg/labels"
	"code.forgejo.org/forgejo/runner/v11/testutils"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestRunDaemonGracefulShutdown(t *testing.T) {
	// Key assertions for graceful shutdown test:
	//
	// - ctx passed to createRunner, createPoller, and Shutdown must outlive signalContext passed to runDaemon, allowing
	// the poller to operate without errors after termination signal is received: #1
	//
	// - When shutting down, the order of operations should be: close signalContext, which causes Shutdown mock to be
	// invoked, and Shutdown mock causes the Poll method to be stopped: #2

	mockClient := mock_client.NewClient(t)
	mockRunner := mock_runner.NewRunnerInterface(t)
	mockPoller := mock_poller.NewPoller(t)

	defer testutils.MockVariable(&initializeConfig, func(configFile *string) (*config.Config, error) {
		return &config.Config{
			Runner: config.Runner{
				// Default ShutdownTimeout of 0s won't work for the graceful shutdown test.
				ShutdownTimeout: 30 * time.Second,
			},
		}, nil
	})()
	defer testutils.MockVariable(&initLogging, func(cfg *config.Config) {})()
	defer testutils.MockVariable(&loadRegistration, func(cfg *config.Config) (*config.Registration, error) {
		return &config.Registration{}, nil
	})()
	defer testutils.MockVariable(&extractLabels, func(cfg *config.Config, reg *config.Registration) labels.Labels {
		return labels.Labels{}
	})()
	defer testutils.MockVariable(&configCheck, func(ctx context.Context, cfg *config.Config, ls labels.Labels) error {
		return nil
	})()
	defer testutils.MockVariable(&createClient, func(cfg *config.Config, reg *config.Registration) client.Client {
		return mockClient
	})()
	var runnerContext context.Context
	defer testutils.MockVariable(&createRunner, func(ctx context.Context, cfg *config.Config, reg *config.Registration, cli client.Client, ls labels.Labels) (run.RunnerInterface, string, error) {
		runnerContext = ctx
		return mockRunner, "runner", nil
	})()
	var pollerContext context.Context
	defer testutils.MockVariable(&createPoller, func(ctx context.Context, cfg *config.Config, cli client.Client, runner run.RunnerInterface) poll.Poller {
		pollerContext = ctx
		return mockPoller
	})()

	pollBegunChannel := make(chan interface{})
	shutdownChannel := make(chan interface{})
	mockPoller.On("Poll").Run(func(args mock.Arguments) {
		close(pollBegunChannel)
		// Simulate running the poll by waiting and doing nothing until shutdownChannel says Shutdown was invoked
		require.NotNil(t, pollerContext)
		select {
		case <-pollerContext.Done():
			assert.Fail(t, "pollerContext was closed before shutdownChannel") // #1
			return
		case <-shutdownChannel:
			return
		}
	})
	mockPoller.On("Shutdown", mock.Anything).Run(func(args mock.Arguments) {
		shutdownContext := args.Get(0).(context.Context)
		select {
		case <-shutdownContext.Done():
			assert.Fail(t, "shutdownContext was closed, but was expected to be open") // #1
			return
		case <-runnerContext.Done():
			assert.Fail(t, "runnerContext was closed, but was expected to be open") // #1
			return
		case <-time.After(time.Microsecond):
			close(shutdownChannel)
			return
		}
	}).Return(nil)

	// When runDaemon is begun, it will run "forever" until the passed-in context is done.  So, let's start that in a goroutine...
	mockSignalContext, cancelSignal := context.WithCancel(t.Context())
	runDaemonComplete := make(chan interface{})
	go func() {
		configFile := "config.yaml"
		err := runDaemon(mockSignalContext, &configFile)
		close(runDaemonComplete)
		require.NoError(t, err)
	}()

	// Wait until runDaemon reaches poller.Poll(), where we expect graceful shutdown to trigger
	<-pollBegunChannel

	// Now we'll signal to the daemon to begin graceful shutdown; this begins the events described in #2
	cancelSignal()

	// Wait for the daemon goroutine to stop
	<-runDaemonComplete
}