agent/agenttest/agent.go

Rama main · réplica de lectura del repositorio original.

package agenttest

import (
	"net/url"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/coder/coder/v2/agent"
	"github.com/coder/coder/v2/agent/agentcontextconfig"
	"github.com/coder/coder/v2/codersdk/agentsdk"
	"github.com/coder/coder/v2/testutil"
)

// New starts a new agent for use in tests.
// The agent will use the provided coder URL and session token.
// The options passed to agent.New() can be modified by passing an optional
// variadic func(*agent.Options).
// Returns the agent. Closing the agent is handled by the test cleanup.
// It is the responsibility of the caller to call coderdtest.AwaitWorkspaceAgents
// to ensure agent is connected.
func New(t testing.TB, coderURL *url.URL, agentToken string, opts ...func(*agent.Options)) agent.Agent {
	t.Helper()

	var o agent.Options
	log := testutil.Logger(t).Named("agent")
	o.Logger = log
	o.SocketPath = testutil.AgentSocketPath(t)

	for _, opt := range opts {
		opt(&o)
	}

	if o.Client == nil {
		agentClient := agentsdk.New(coderURL, agentsdk.WithFixedToken(agentToken))
		agentClient.SDK.SetLogger(log)
		o.Client = agentClient
	}

	if o.LogDir == "" {
		o.LogDir = t.TempDir()
	}

	agt := agent.New(o)
	t.Cleanup(func() {
		assert.NoError(t, agt.Close(), "failed to close agent during cleanup")
	})

	return agt
}

// WithContextConfigFromEnv returns an agent option that
// populates ContextConfig from the current environment.
func WithContextConfigFromEnv() func(*agent.Options) {
	return func(o *agent.Options) {
		o.ContextConfig = agentcontextconfig.ReadEnvConfig()
	}
}