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.

codersdk/workspacesdk/agentconn.go

@@ -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)
 }

codersdk/workspacesdk/agentconnmock/agentconnmock.go

@@ -0,0 +1,373 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: .. (interfaces: AgentConn)
+//
+// Generated by this command:
+//
+//	mockgen -destination ./agentconnmock.go -package agentconnmock .. AgentConn
+//
+
+// Package agentconnmock is a generated GoMock package.
+package agentconnmock
+
+import (
+	context "context"
+	io "io"
+	net "net"
+	reflect "reflect"
+	time "time"
+
+	slog "cdr.dev/slog"
+	codersdk "github.com/coder/coder/v2/codersdk"
+	healthsdk "github.com/coder/coder/v2/codersdk/healthsdk"
+	workspacesdk "github.com/coder/coder/v2/codersdk/workspacesdk"
+	tailnet "github.com/coder/coder/v2/tailnet"
+	uuid "github.com/google/uuid"
+	gomock "go.uber.org/mock/gomock"
+	ssh "golang.org/x/crypto/ssh"
+	gonet "gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+	ipnstate "tailscale.com/ipn/ipnstate"
+	speedtest "tailscale.com/net/speedtest"
+)
+
+// MockAgentConn is a mock of AgentConn interface.
+type MockAgentConn struct {
+	ctrl     *gomock.Controller
+	recorder *MockAgentConnMockRecorder
+	isgomock struct{}
+}
+
+// MockAgentConnMockRecorder is the mock recorder for MockAgentConn.
+type MockAgentConnMockRecorder struct {
+	mock *MockAgentConn
+}
+
+// NewMockAgentConn creates a new mock instance.
+func NewMockAgentConn(ctrl *gomock.Controller) *MockAgentConn {
+	mock := &MockAgentConn{ctrl: ctrl}
+	mock.recorder = &MockAgentConnMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockAgentConn) EXPECT() *MockAgentConnMockRecorder {
+	return m.recorder
+}
+
+// AwaitReachable mocks base method.
+func (m *MockAgentConn) AwaitReachable(ctx context.Context) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "AwaitReachable", ctx)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// AwaitReachable indicates an expected call of AwaitReachable.
+func (mr *MockAgentConnMockRecorder) AwaitReachable(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AwaitReachable", reflect.TypeOf((*MockAgentConn)(nil).AwaitReachable), ctx)
+}
+
+// Close mocks base method.
+func (m *MockAgentConn) Close() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Close")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Close indicates an expected call of Close.
+func (mr *MockAgentConnMockRecorder) Close() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockAgentConn)(nil).Close))
+}
+
+// DebugLogs mocks base method.
+func (m *MockAgentConn) DebugLogs(ctx context.Context) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DebugLogs", ctx)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DebugLogs indicates an expected call of DebugLogs.
+func (mr *MockAgentConnMockRecorder) DebugLogs(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DebugLogs", reflect.TypeOf((*MockAgentConn)(nil).DebugLogs), ctx)
+}
+
+// DebugMagicsock mocks base method.
+func (m *MockAgentConn) DebugMagicsock(ctx context.Context) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DebugMagicsock", ctx)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DebugMagicsock indicates an expected call of DebugMagicsock.
+func (mr *MockAgentConnMockRecorder) DebugMagicsock(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DebugMagicsock", reflect.TypeOf((*MockAgentConn)(nil).DebugMagicsock), ctx)
+}
+
+// DebugManifest mocks base method.
+func (m *MockAgentConn) DebugManifest(ctx context.Context) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DebugManifest", ctx)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DebugManifest indicates an expected call of DebugManifest.
+func (mr *MockAgentConnMockRecorder) DebugManifest(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DebugManifest", reflect.TypeOf((*MockAgentConn)(nil).DebugManifest), ctx)
+}
+
+// DialContext mocks base method.
+func (m *MockAgentConn) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DialContext", ctx, network, addr)
+	ret0, _ := ret[0].(net.Conn)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DialContext indicates an expected call of DialContext.
+func (mr *MockAgentConnMockRecorder) DialContext(ctx, network, addr any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DialContext", reflect.TypeOf((*MockAgentConn)(nil).DialContext), ctx, network, addr)
+}
+
+// GetPeerDiagnostics mocks base method.
+func (m *MockAgentConn) GetPeerDiagnostics() tailnet.PeerDiagnostics {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPeerDiagnostics")
+	ret0, _ := ret[0].(tailnet.PeerDiagnostics)
+	return ret0
+}
+
+// GetPeerDiagnostics indicates an expected call of GetPeerDiagnostics.
+func (mr *MockAgentConnMockRecorder) GetPeerDiagnostics() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPeerDiagnostics", reflect.TypeOf((*MockAgentConn)(nil).GetPeerDiagnostics))
+}
+
+// ListContainers mocks base method.
+func (m *MockAgentConn) ListContainers(ctx context.Context) (codersdk.WorkspaceAgentListContainersResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ListContainers", ctx)
+	ret0, _ := ret[0].(codersdk.WorkspaceAgentListContainersResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ListContainers indicates an expected call of ListContainers.
+func (mr *MockAgentConnMockRecorder) ListContainers(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListContainers", reflect.TypeOf((*MockAgentConn)(nil).ListContainers), ctx)
+}
+
+// ListeningPorts mocks base method.
+func (m *MockAgentConn) ListeningPorts(ctx context.Context) (codersdk.WorkspaceAgentListeningPortsResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ListeningPorts", ctx)
+	ret0, _ := ret[0].(codersdk.WorkspaceAgentListeningPortsResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ListeningPorts indicates an expected call of ListeningPorts.
+func (mr *MockAgentConnMockRecorder) ListeningPorts(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListeningPorts", reflect.TypeOf((*MockAgentConn)(nil).ListeningPorts), ctx)
+}
+
+// Netcheck mocks base method.
+func (m *MockAgentConn) Netcheck(ctx context.Context) (healthsdk.AgentNetcheckReport, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Netcheck", ctx)
+	ret0, _ := ret[0].(healthsdk.AgentNetcheckReport)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Netcheck indicates an expected call of Netcheck.
+func (mr *MockAgentConnMockRecorder) Netcheck(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Netcheck", reflect.TypeOf((*MockAgentConn)(nil).Netcheck), ctx)
+}
+
+// Ping mocks base method.
+func (m *MockAgentConn) Ping(ctx context.Context) (time.Duration, bool, *ipnstate.PingResult, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Ping", ctx)
+	ret0, _ := ret[0].(time.Duration)
+	ret1, _ := ret[1].(bool)
+	ret2, _ := ret[2].(*ipnstate.PingResult)
+	ret3, _ := ret[3].(error)
+	return ret0, ret1, ret2, ret3
+}
+
+// Ping indicates an expected call of Ping.
+func (mr *MockAgentConnMockRecorder) Ping(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ping", reflect.TypeOf((*MockAgentConn)(nil).Ping), ctx)
+}
+
+// PrometheusMetrics mocks base method.
+func (m *MockAgentConn) PrometheusMetrics(ctx context.Context) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PrometheusMetrics", ctx)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PrometheusMetrics indicates an expected call of PrometheusMetrics.
+func (mr *MockAgentConnMockRecorder) PrometheusMetrics(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PrometheusMetrics", reflect.TypeOf((*MockAgentConn)(nil).PrometheusMetrics), ctx)
+}
+
+// ReconnectingPTY mocks base method.
+func (m *MockAgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string, initOpts ...workspacesdk.AgentReconnectingPTYInitOption) (net.Conn, error) {
+	m.ctrl.T.Helper()
+	varargs := []any{ctx, id, height, width, command}
+	for _, a := range initOpts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "ReconnectingPTY", varargs...)
+	ret0, _ := ret[0].(net.Conn)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ReconnectingPTY indicates an expected call of ReconnectingPTY.
+func (mr *MockAgentConnMockRecorder) ReconnectingPTY(ctx, id, height, width, command any, initOpts ...any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]any{ctx, id, height, width, command}, initOpts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReconnectingPTY", reflect.TypeOf((*MockAgentConn)(nil).ReconnectingPTY), varargs...)
+}
+
+// RecreateDevcontainer mocks base method.
+func (m *MockAgentConn) RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RecreateDevcontainer", ctx, devcontainerID)
+	ret0, _ := ret[0].(codersdk.Response)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// RecreateDevcontainer indicates an expected call of RecreateDevcontainer.
+func (mr *MockAgentConnMockRecorder) RecreateDevcontainer(ctx, devcontainerID any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecreateDevcontainer", reflect.TypeOf((*MockAgentConn)(nil).RecreateDevcontainer), ctx, devcontainerID)
+}
+
+// SSH mocks base method.
+func (m *MockAgentConn) SSH(ctx context.Context) (*gonet.TCPConn, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SSH", ctx)
+	ret0, _ := ret[0].(*gonet.TCPConn)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SSH indicates an expected call of SSH.
+func (mr *MockAgentConnMockRecorder) SSH(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSH", reflect.TypeOf((*MockAgentConn)(nil).SSH), ctx)
+}
+
+// SSHClient mocks base method.
+func (m *MockAgentConn) SSHClient(ctx context.Context) (*ssh.Client, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SSHClient", ctx)
+	ret0, _ := ret[0].(*ssh.Client)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SSHClient indicates an expected call of SSHClient.
+func (mr *MockAgentConnMockRecorder) SSHClient(ctx any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSHClient", reflect.TypeOf((*MockAgentConn)(nil).SSHClient), ctx)
+}
+
+// SSHClientOnPort mocks base method.
+func (m *MockAgentConn) SSHClientOnPort(ctx context.Context, port uint16) (*ssh.Client, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SSHClientOnPort", ctx, port)
+	ret0, _ := ret[0].(*ssh.Client)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SSHClientOnPort indicates an expected call of SSHClientOnPort.
+func (mr *MockAgentConnMockRecorder) SSHClientOnPort(ctx, port any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSHClientOnPort", reflect.TypeOf((*MockAgentConn)(nil).SSHClientOnPort), ctx, port)
+}
+
+// SSHOnPort mocks base method.
+func (m *MockAgentConn) SSHOnPort(ctx context.Context, port uint16) (*gonet.TCPConn, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SSHOnPort", ctx, port)
+	ret0, _ := ret[0].(*gonet.TCPConn)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SSHOnPort indicates an expected call of SSHOnPort.
+func (mr *MockAgentConnMockRecorder) SSHOnPort(ctx, port any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SSHOnPort", reflect.TypeOf((*MockAgentConn)(nil).SSHOnPort), ctx, port)
+}
+
+// Speedtest mocks base method.
+func (m *MockAgentConn) Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Speedtest", ctx, direction, duration)
+	ret0, _ := ret[0].([]speedtest.Result)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Speedtest indicates an expected call of Speedtest.
+func (mr *MockAgentConnMockRecorder) Speedtest(ctx, direction, duration any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Speedtest", reflect.TypeOf((*MockAgentConn)(nil).Speedtest), ctx, direction, duration)
+}
+
+// TailnetConn mocks base method.
+func (m *MockAgentConn) TailnetConn() *tailnet.Conn {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "TailnetConn")
+	ret0, _ := ret[0].(*tailnet.Conn)
+	return ret0
+}
+
+// TailnetConn indicates an expected call of TailnetConn.
+func (mr *MockAgentConnMockRecorder) TailnetConn() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TailnetConn", reflect.TypeOf((*MockAgentConn)(nil).TailnetConn))
+}
+
+// WatchContainers mocks base method.
+func (m *MockAgentConn) WatchContainers(ctx context.Context, logger slog.Logger) (<-chan codersdk.WorkspaceAgentListContainersResponse, io.Closer, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WatchContainers", ctx, logger)
+	ret0, _ := ret[0].(<-chan codersdk.WorkspaceAgentListContainersResponse)
+	ret1, _ := ret[1].(io.Closer)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
+}
+
+// WatchContainers indicates an expected call of WatchContainers.
+func (mr *MockAgentConnMockRecorder) WatchContainers(ctx, logger any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WatchContainers", reflect.TypeOf((*MockAgentConn)(nil).WatchContainers), ctx, logger)
+}

codersdk/workspacesdk/agentconnmock/doc.go

@@ -0,0 +1,4 @@
+// Package agentconnmock contains a mock implementation of workspacesdk.AgentConn for use in tests.
+package agentconnmock
+
+//go:generate mockgen -destination ./agentconnmock.go -package agentconnmock .. AgentConn

codersdk/workspacesdk/workspacesdk.go

@@ -202,7 +202,7 @@ func (c *Client) RewriteDERPMap(derpMap *tailcfg.DERPMap) {
 	tailnet.RewriteDERPMapDefaultRelay(context.Background(), c.client.Logger(), derpMap, c.client.URL)
 }
 
-func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *DialAgentOptions) (agentConn *AgentConn, err error) {
+func (c *Client) DialAgent(dialCtx context.Context, agentID uuid.UUID, options *DialAgentOptions) (agentConn AgentConn, err error) {
 	if options == nil {
 		options = &DialAgentOptions{}
 	}