refactor: convert workspacesdk.AgentConn to an interface (#19392)

Fixes https://github.com/coder/internal/issues/907

We convert `workspacesdk.AgentConn` to an interface and generate a mock
for it. This allows writing `coderd` tests that rely on the agent's HTTP
api to not have to set up an entire tailnet networking stack.
This commit is contained in:
Danielle Maywood
2025-08-20 10:00:44 +01:00
committed by GitHub
parent 23c494f36b
commit 5e84d257b7
18 changed files with 667 additions and 143 deletions
+57 -26
View File
@@ -34,8 +34,8 @@ import (
// to the WorkspaceAgentConn, or it may be shared in the case of coderd. If the
// conn is shared and closing it is undesirable, you may return ErrNoClose from
// opts.CloseFunc. This will ensure the underlying conn is not closed.
func NewAgentConn(conn *tailnet.Conn, opts AgentConnOptions) *AgentConn {
return &AgentConn{
func NewAgentConn(conn *tailnet.Conn, opts AgentConnOptions) AgentConn {
return &agentConn{
Conn: conn,
opts: opts,
}
@@ -43,23 +43,54 @@ func NewAgentConn(conn *tailnet.Conn, opts AgentConnOptions) *AgentConn {
// AgentConn represents a connection to a workspace agent.
// @typescript-ignore AgentConn
type AgentConn struct {
type AgentConn interface {
TailnetConn() *tailnet.Conn
AwaitReachable(ctx context.Context) bool
Close() error
DebugLogs(ctx context.Context) ([]byte, error)
DebugMagicsock(ctx context.Context) ([]byte, error)
DebugManifest(ctx context.Context) ([]byte, error)
DialContext(ctx context.Context, network string, addr string) (net.Conn, error)
GetPeerDiagnostics() tailnet.PeerDiagnostics
ListContainers(ctx context.Context) (codersdk.WorkspaceAgentListContainersResponse, error)
ListeningPorts(ctx context.Context) (codersdk.WorkspaceAgentListeningPortsResponse, error)
Netcheck(ctx context.Context) (healthsdk.AgentNetcheckReport, error)
Ping(ctx context.Context) (time.Duration, bool, *ipnstate.PingResult, error)
PrometheusMetrics(ctx context.Context) ([]byte, error)
ReconnectingPTY(ctx context.Context, id uuid.UUID, height uint16, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error)
RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error)
SSH(ctx context.Context) (*gonet.TCPConn, error)
SSHClient(ctx context.Context) (*ssh.Client, error)
SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error)
SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error)
Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error)
WatchContainers(ctx context.Context, logger slog.Logger) (<-chan codersdk.WorkspaceAgentListContainersResponse, io.Closer, error)
}
// AgentConn represents a connection to a workspace agent.
// @typescript-ignore AgentConn
type agentConn struct {
*tailnet.Conn
opts AgentConnOptions
}
func (c *agentConn) TailnetConn() *tailnet.Conn {
return c.Conn
}
// @typescript-ignore AgentConnOptions
type AgentConnOptions struct {
AgentID uuid.UUID
CloseFunc func() error
}
func (c *AgentConn) agentAddress() netip.Addr {
func (c *agentConn) agentAddress() netip.Addr {
return tailnet.TailscaleServicePrefix.AddrFromUUID(c.opts.AgentID)
}
// AwaitReachable waits for the agent to be reachable.
func (c *AgentConn) AwaitReachable(ctx context.Context) bool {
func (c *agentConn) AwaitReachable(ctx context.Context) bool {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -68,7 +99,7 @@ func (c *AgentConn) AwaitReachable(ctx context.Context) bool {
// Ping pings the agent and returns the round-trip time.
// The bool returns true if the ping was made P2P.
func (c *AgentConn) Ping(ctx context.Context) (time.Duration, bool, *ipnstate.PingResult, error) {
func (c *agentConn) Ping(ctx context.Context) (time.Duration, bool, *ipnstate.PingResult, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -76,7 +107,7 @@ func (c *AgentConn) Ping(ctx context.Context) (time.Duration, bool, *ipnstate.Pi
}
// Close ends the connection to the workspace agent.
func (c *AgentConn) Close() error {
func (c *agentConn) Close() error {
var cerr error
if c.opts.CloseFunc != nil {
cerr = c.opts.CloseFunc()
@@ -131,7 +162,7 @@ type ReconnectingPTYRequest struct {
// ReconnectingPTY spawns a new reconnecting terminal session.
// `ReconnectingPTYRequest` should be JSON marshaled and written to the returned net.Conn.
// Raw terminal output will be read from the returned net.Conn.
func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error) {
func (c *agentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string, initOpts ...AgentReconnectingPTYInitOption) (net.Conn, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -171,13 +202,13 @@ func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, w
// SSH pipes the SSH protocol over the returned net.Conn.
// This connects to the built-in SSH server in the workspace agent.
func (c *AgentConn) SSH(ctx context.Context) (*gonet.TCPConn, error) {
func (c *agentConn) SSH(ctx context.Context) (*gonet.TCPConn, error) {
return c.SSHOnPort(ctx, AgentSSHPort)
}
// SSHOnPort pipes the SSH protocol over the returned net.Conn.
// This connects to the built-in SSH server in the workspace agent on the specified port.
func (c *AgentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error) {
func (c *agentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -190,12 +221,12 @@ func (c *AgentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn,
}
// SSHClient calls SSH to create a client
func (c *AgentConn) SSHClient(ctx context.Context) (*ssh.Client, error) {
func (c *agentConn) SSHClient(ctx context.Context) (*ssh.Client, error) {
return c.SSHClientOnPort(ctx, AgentSSHPort)
}
// SSHClientOnPort calls SSH to create a client on a specific port
func (c *AgentConn) SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error) {
func (c *agentConn) SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -218,7 +249,7 @@ func (c *AgentConn) SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Clie
}
// Speedtest runs a speedtest against the workspace agent.
func (c *AgentConn) Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
func (c *agentConn) Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -242,7 +273,7 @@ func (c *AgentConn) Speedtest(ctx context.Context, direction speedtest.Direction
// DialContext dials the address provided in the workspace agent.
// The network must be "tcp" or "udp".
func (c *AgentConn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
func (c *agentConn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -265,7 +296,7 @@ func (c *AgentConn) DialContext(ctx context.Context, network string, addr string
}
// ListeningPorts lists the ports that are currently in use by the workspace.
func (c *AgentConn) ListeningPorts(ctx context.Context) (codersdk.WorkspaceAgentListeningPortsResponse, error) {
func (c *agentConn) ListeningPorts(ctx context.Context) (codersdk.WorkspaceAgentListeningPortsResponse, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/api/v0/listening-ports", nil)
@@ -282,7 +313,7 @@ func (c *AgentConn) ListeningPorts(ctx context.Context) (codersdk.WorkspaceAgent
}
// Netcheck returns a network check report from the workspace agent.
func (c *AgentConn) Netcheck(ctx context.Context) (healthsdk.AgentNetcheckReport, error) {
func (c *agentConn) Netcheck(ctx context.Context) (healthsdk.AgentNetcheckReport, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/api/v0/netcheck", nil)
@@ -299,7 +330,7 @@ func (c *AgentConn) Netcheck(ctx context.Context) (healthsdk.AgentNetcheckReport
}
// DebugMagicsock makes a request to the workspace agent's magicsock debug endpoint.
func (c *AgentConn) DebugMagicsock(ctx context.Context) ([]byte, error) {
func (c *agentConn) DebugMagicsock(ctx context.Context) ([]byte, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/debug/magicsock", nil)
@@ -319,7 +350,7 @@ func (c *AgentConn) DebugMagicsock(ctx context.Context) ([]byte, error) {
// DebugManifest returns the agent's in-memory manifest. Unfortunately this must
// be returns as a []byte to avoid an import cycle.
func (c *AgentConn) DebugManifest(ctx context.Context) ([]byte, error) {
func (c *agentConn) DebugManifest(ctx context.Context) ([]byte, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/debug/manifest", nil)
@@ -338,7 +369,7 @@ func (c *AgentConn) DebugManifest(ctx context.Context) ([]byte, error) {
}
// DebugLogs returns up to the last 10MB of `/tmp/coder-agent.log`
func (c *AgentConn) DebugLogs(ctx context.Context) ([]byte, error) {
func (c *agentConn) DebugLogs(ctx context.Context) ([]byte, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/debug/logs", nil)
@@ -357,7 +388,7 @@ func (c *AgentConn) DebugLogs(ctx context.Context) ([]byte, error) {
}
// PrometheusMetrics returns a response from the agent's prometheus metrics endpoint
func (c *AgentConn) PrometheusMetrics(ctx context.Context) ([]byte, error) {
func (c *agentConn) PrometheusMetrics(ctx context.Context) ([]byte, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/debug/prometheus", nil)
@@ -376,7 +407,7 @@ func (c *AgentConn) PrometheusMetrics(ctx context.Context) ([]byte, error) {
}
// ListContainers returns a response from the agent's containers endpoint
func (c *AgentConn) ListContainers(ctx context.Context) (codersdk.WorkspaceAgentListContainersResponse, error) {
func (c *agentConn) ListContainers(ctx context.Context) (codersdk.WorkspaceAgentListContainersResponse, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodGet, "/api/v0/containers", nil)
@@ -391,7 +422,7 @@ func (c *AgentConn) ListContainers(ctx context.Context) (codersdk.WorkspaceAgent
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
func (c *AgentConn) WatchContainers(ctx context.Context, logger slog.Logger) (<-chan codersdk.WorkspaceAgentListContainersResponse, io.Closer, error) {
func (c *agentConn) WatchContainers(ctx context.Context, logger slog.Logger) (<-chan codersdk.WorkspaceAgentListContainersResponse, io.Closer, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -427,7 +458,7 @@ func (c *AgentConn) WatchContainers(ctx context.Context, logger slog.Logger) (<-
// RecreateDevcontainer recreates a devcontainer with the given container.
// This is a blocking call and will wait for the container to be recreated.
func (c *AgentConn) RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error) {
func (c *agentConn) RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/containers/devcontainers/"+devcontainerID+"/recreate", nil)
@@ -446,7 +477,7 @@ func (c *AgentConn) RecreateDevcontainer(ctx context.Context, devcontainerID str
}
// apiRequest makes a request to the workspace agent's HTTP API server.
func (c *AgentConn) apiRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
func (c *agentConn) apiRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
ctx, span := tracing.StartSpan(ctx)
defer span.End()
@@ -463,7 +494,7 @@ func (c *AgentConn) apiRequest(ctx context.Context, method, path string, body io
// apiClient returns an HTTP client that can be used to make
// requests to the workspace agent's HTTP API server.
func (c *AgentConn) apiClient() *http.Client {
func (c *agentConn) apiClient() *http.Client {
return &http.Client{
Transport: &http.Transport{
// Disable keep alives as we're usually only making a single
@@ -504,6 +535,6 @@ func (c *AgentConn) apiClient() *http.Client {
}
}
func (c *AgentConn) GetPeerDiagnostics() tailnet.PeerDiagnostics {
func (c *agentConn) GetPeerDiagnostics() tailnet.PeerDiagnostics {
return c.Conn.GetPeerDiagnostics(c.opts.AgentID)
}