Skip to content

Commit ac0d625

Browse files
coadlercoadler
committed
fixup! tests
1 parent d484a67 commit ac0d625

File tree

3 files changed

+53
-194
lines changed

3 files changed

+53
-194
lines changed

‎agent/agent_test.go

Lines changed: 46 additions & 187 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ import (
3434
"github.com/stretchr/testify/require"
3535
"go.uber.org/goleak"
3636
"golang.org/x/crypto/ssh"
37-
"golang.org/x/exp/maps"
3837
"golang.org/x/xerrors"
3938
"tailscale.com/net/speedtest"
4039
"tailscale.com/tailcfg"
@@ -882,16 +881,15 @@ func TestAgent_StartupScript(t *testing.T) {
882881
t.Run("Success", func(t *testing.T) {
883882
t.Parallel()
884883
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
885-
client := &client{
886-
t: t,
887-
agentID: uuid.New(),
888-
manifest: agentsdk.Manifest{
884+
client := agenttest.NewClient(t,
885+
uuid.New(),
886+
agentsdk.Manifest{
889887
StartupScript: command,
890888
DERPMap: &tailcfg.DERPMap{},
891889
},
892-
statsChan: make(chan *agentsdk.Stats),
893-
coordinator: tailnet.NewCoordinator(logger),
894-
}
890+
make(chan *agentsdk.Stats),
891+
tailnet.NewCoordinator(logger),
892+
)
895893
closer := agent.New(agent.Options{
896894
Client: client,
897895
Filesystem: afero.NewMemMapFs(),
@@ -902,36 +900,35 @@ func TestAgent_StartupScript(t *testing.T) {
902900
_ = closer.Close()
903901
})
904902
assert.Eventually(t, func() bool {
905-
got := client.getLifecycleStates()
903+
got := client.GetLifecycleStates()
906904
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
907905
}, testutil.WaitShort, testutil.IntervalMedium)
908906

909-
require.Len(t, client.getStartupLogs(), 1)
910-
require.Equal(t, output, client.getStartupLogs()[0].Output)
907+
require.Len(t, client.GetStartupLogs(), 1)
908+
require.Equal(t, output, client.GetStartupLogs()[0].Output)
911909
})
912910
// This ensures that even when coderd sends back that the startup
913911
// script has written too many lines it will still succeed!
914912
t.Run("OverflowsAndSkips", func(t *testing.T) {
915913
t.Parallel()
916914
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
917-
client := &client{
918-
t: t,
919-
agentID: uuid.New(),
920-
manifest: agentsdk.Manifest{
915+
client := agenttest.NewClient(t,
916+
uuid.New(),
917+
agentsdk.Manifest{
921918
StartupScript: command,
922919
DERPMap: &tailcfg.DERPMap{},
923920
},
924-
patchWorkspaceLogs: func() error {
925-
resp := httptest.NewRecorder()
926-
httpapi.Write(context.Background(), resp, http.StatusRequestEntityTooLarge, codersdk.Response{
927-
Message: "Too many lines!",
928-
})
929-
res := resp.Result()
930-
defer res.Body.Close()
931-
return codersdk.ReadBodyAsError(res)
932-
},
933-
statsChan: make(chan *agentsdk.Stats),
934-
coordinator: tailnet.NewCoordinator(logger),
921+
make(chan *agentsdk.Stats, 100),
922+
tailnet.NewCoordinator(logger),
923+
)
924+
client.PatchWorkspaceLogs = func() error {
925+
resp := httptest.NewRecorder()
926+
httpapi.Write(context.Background(), resp, http.StatusRequestEntityTooLarge, codersdk.Response{
927+
Message: "Too many lines!",
928+
})
929+
res := resp.Result()
930+
defer res.Body.Close()
931+
return codersdk.ReadBodyAsError(res)
935932
}
936933
closer := agent.New(agent.Options{
937934
Client: client,
@@ -943,10 +940,10 @@ func TestAgent_StartupScript(t *testing.T) {
943940
_ = closer.Close()
944941
})
945942
assert.Eventually(t, func() bool {
946-
got := client.getLifecycleStates()
943+
got := client.GetLifecycleStates()
947944
return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
948945
}, testutil.WaitShort, testutil.IntervalMedium)
949-
require.Len(t, client.getStartupLogs(), 0)
946+
require.Len(t, client.GetStartupLogs(), 0)
950947
})
951948
}
952949

@@ -1283,17 +1280,17 @@ func TestAgent_Lifecycle(t *testing.T) {
12831280
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
12841281
expected := "this-is-shutdown"
12851282
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
1286-
client := &client{
1287-
t: t,
1288-
agentID: uuid.New(),
1289-
manifest: agentsdk.Manifest{
1283+
1284+
client := agenttest.NewClient(t,
1285+
uuid.New(),
1286+
agentsdk.Manifest{
12901287
DERPMap: derpMap,
12911288
StartupScript: "echo 1",
12921289
ShutdownScript: "echo " + expected,
12931290
},
1294-
statsChan: make(chan *agentsdk.Stats),
1295-
coordinator: tailnet.NewCoordinator(logger),
1296-
}
1291+
make(chan *agentsdk.Stats, 100),
1292+
tailnet.NewCoordinator(logger),
1293+
)
12971294

12981295
fs := afero.NewMemMapFs()
12991296
agent := agent.New(agent.Options{
@@ -1554,17 +1551,16 @@ func TestAgent_Reconnect(t *testing.T) {
15541551
defer coordinator.Close()
15551552

15561553
agentID := uuid.New()
1557-
statsCh := make(chan *agentsdk.Stats)
1554+
statsCh := make(chan *agentsdk.Stats, 50)
15581555
derpMap, _ := tailnettest.RunDERPAndSTUN(t)
1559-
client := &client{
1560-
t: t,
1561-
agentID: agentID,
1562-
manifest: agentsdk.Manifest{
1556+
client := agenttest.NewClient(t,
1557+
agentID,
1558+
agentsdk.Manifest{
15631559
DERPMap: derpMap,
15641560
},
1565-
statsChan: statsCh,
1566-
coordinator: coordinator,
1567-
}
1561+
statsCh,
1562+
coordinator,
1563+
)
15681564
initialized := atomic.Int32{}
15691565
closer := agent.New(agent.Options{
15701566
ExchangeToken: func(ctx context.Context) (string, error) {
@@ -1579,7 +1575,7 @@ func TestAgent_Reconnect(t *testing.T) {
15791575
require.Eventually(t, func() bool {
15801576
return coordinator.Node(agentID) != nil
15811577
}, testutil.WaitShort, testutil.IntervalFast)
1582-
client.lastWorkspaceAgent()
1578+
client.LastWorkspaceAgent()
15831579
require.Eventually(t, func() bool {
15841580
return initialized.Load() == 2
15851581
}, testutil.WaitShort, testutil.IntervalFast)
@@ -1591,16 +1587,15 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
15911587
coordinator := tailnet.NewCoordinator(logger)
15921588
defer coordinator.Close()
15931589

1594-
client := &client{
1595-
t: t,
1596-
agentID: uuid.New(),
1597-
manifest: agentsdk.Manifest{
1590+
client := agenttest.NewClient(t,
1591+
uuid.New(),
1592+
agentsdk.Manifest{
15981593
GitAuthConfigs: 1,
15991594
DERPMap: &tailcfg.DERPMap{},
16001595
},
1601-
statsChan: make(chan *agentsdk.Stats),
1602-
coordinator: coordinator,
1603-
}
1596+
make(chan *agentsdk.Stats, 50),
1597+
coordinator,
1598+
)
16041599
filesystem := afero.NewMemMapFs()
16051600
closer := agent.New(agent.Options{
16061601
ExchangeToken: func(ctx context.Context) (string, error) {
@@ -1685,12 +1680,6 @@ func setupSSHSession(t *testing.T, options agentsdk.Manifest) *ssh.Session {
16851680
return session
16861681
}
16871682

1688-
type closeFunc func() error
1689-
1690-
func (c closeFunc) Close() error {
1691-
return c()
1692-
}
1693-
16941683
func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Duration, opts ...func(agent.Options) agent.Options) (
16951684
*codersdk.WorkspaceAgentConn,
16961685
*agenttest.Client,
@@ -1803,136 +1792,6 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
18031792
assert.Equal(t, len(payload), n, "payload length does not match")
18041793
}
18051794

1806-
type client struct {
1807-
t *testing.T
1808-
agentID uuid.UUID
1809-
manifest agentsdk.Manifest
1810-
metadata map[string]agentsdk.PostMetadataRequest
1811-
statsChan chan *agentsdk.Stats
1812-
coordinator tailnet.Coordinator
1813-
lastWorkspaceAgent func()
1814-
patchWorkspaceLogs func() error
1815-
1816-
mu sync.Mutex // Protects following.
1817-
lifecycleStates []codersdk.WorkspaceAgentLifecycle
1818-
startup agentsdk.PostStartupRequest
1819-
logs []agentsdk.StartupLog
1820-
}
1821-
1822-
func (c *client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
1823-
return c.manifest, nil
1824-
}
1825-
1826-
func (c *client) Listen(_ context.Context) (net.Conn, error) {
1827-
clientConn, serverConn := net.Pipe()
1828-
closed := make(chan struct{})
1829-
c.lastWorkspaceAgent = func() {
1830-
_ = serverConn.Close()
1831-
_ = clientConn.Close()
1832-
<-closed
1833-
}
1834-
c.t.Cleanup(c.lastWorkspaceAgent)
1835-
go func() {
1836-
_ = c.coordinator.ServeAgent(serverConn, c.agentID, "")
1837-
close(closed)
1838-
}()
1839-
return clientConn, nil
1840-
}
1841-
1842-
func (c *client) ReportStats(ctx context.Context, _ slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error) {
1843-
doneCh := make(chan struct{})
1844-
ctx, cancel := context.WithCancel(ctx)
1845-
1846-
go func() {
1847-
defer close(doneCh)
1848-
1849-
setInterval(500 * time.Millisecond)
1850-
for {
1851-
select {
1852-
case <-ctx.Done():
1853-
return
1854-
case stat := <-statsChan:
1855-
select {
1856-
case c.statsChan <- stat:
1857-
case <-ctx.Done():
1858-
return
1859-
default:
1860-
// We don't want to send old stats.
1861-
continue
1862-
}
1863-
}
1864-
}
1865-
}()
1866-
return closeFunc(func() error {
1867-
cancel()
1868-
<-doneCh
1869-
close(c.statsChan)
1870-
return nil
1871-
}), nil
1872-
}
1873-
1874-
func (c *client) getLifecycleStates() []codersdk.WorkspaceAgentLifecycle {
1875-
c.mu.Lock()
1876-
defer c.mu.Unlock()
1877-
return c.lifecycleStates
1878-
}
1879-
1880-
func (c *client) PostLifecycle(_ context.Context, req agentsdk.PostLifecycleRequest) error {
1881-
c.mu.Lock()
1882-
defer c.mu.Unlock()
1883-
c.lifecycleStates = append(c.lifecycleStates, req.State)
1884-
return nil
1885-
}
1886-
1887-
func (*client) PostAppHealth(_ context.Context, _ agentsdk.PostAppHealthsRequest) error {
1888-
return nil
1889-
}
1890-
1891-
func (c *client) getStartup() agentsdk.PostStartupRequest {
1892-
c.mu.Lock()
1893-
defer c.mu.Unlock()
1894-
return c.startup
1895-
}
1896-
1897-
func (c *client) getMetadata() map[string]agentsdk.PostMetadataRequest {
1898-
c.mu.Lock()
1899-
defer c.mu.Unlock()
1900-
return maps.Clone(c.metadata)
1901-
}
1902-
1903-
func (c *client) PostMetadata(_ context.Context, key string, req agentsdk.PostMetadataRequest) error {
1904-
c.mu.Lock()
1905-
defer c.mu.Unlock()
1906-
if c.metadata == nil {
1907-
c.metadata = make(map[string]agentsdk.PostMetadataRequest)
1908-
}
1909-
c.metadata[key] = req
1910-
return nil
1911-
}
1912-
1913-
func (c *client) PostStartup(_ context.Context, startup agentsdk.PostStartupRequest) error {
1914-
c.mu.Lock()
1915-
defer c.mu.Unlock()
1916-
c.startup = startup
1917-
return nil
1918-
}
1919-
1920-
func (c *client) getStartupLogs() []agentsdk.StartupLog {
1921-
c.mu.Lock()
1922-
defer c.mu.Unlock()
1923-
return c.logs
1924-
}
1925-
1926-
func (c *client) PatchStartupLogs(_ context.Context, logs agentsdk.PatchStartupLogs) error {
1927-
c.mu.Lock()
1928-
defer c.mu.Unlock()
1929-
if c.patchWorkspaceLogs != nil {
1930-
return c.patchWorkspaceLogs()
1931-
}
1932-
c.logs = append(c.logs, logs.Logs...)
1933-
return nil
1934-
}
1935-
19361795
// tempDirUnixSocket returns a temporary directory that can safely hold unix
19371796
// sockets (probably).
19381797
//

‎agent/agenttest/client.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ type Client struct {
3939
metadata map[string]agentsdk.PostMetadataRequest
4040
statsChan chan *agentsdk.Stats
4141
coordinator tailnet.Coordinator
42-
lastWorkspaceAgent func()
43-
patchWorkspaceLogs func() error
42+
LastWorkspaceAgent func()
43+
PatchWorkspaceLogs func() error
4444

4545
mu sync.Mutex // Protects following.
4646
lifecycleStates []codersdk.WorkspaceAgentLifecycle
@@ -55,12 +55,12 @@ func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
5555
func (c *Client) Listen(_ context.Context) (net.Conn, error) {
5656
clientConn, serverConn := net.Pipe()
5757
closed := make(chan struct{})
58-
c.lastWorkspaceAgent = func() {
58+
c.LastWorkspaceAgent = func() {
5959
_ = serverConn.Close()
6060
_ = clientConn.Close()
6161
<-closed
6262
}
63-
c.t.Cleanup(c.lastWorkspaceAgent)
63+
c.t.Cleanup(c.LastWorkspaceAgent)
6464
go func() {
6565
_ = c.coordinator.ServeAgent(serverConn, c.agentID, "")
6666
close(closed)
@@ -155,8 +155,8 @@ func (c *Client) GetStartupLogs() []agentsdk.StartupLog {
155155
func (c *Client) PatchStartupLogs(_ context.Context, logs agentsdk.PatchStartupLogs) error {
156156
c.mu.Lock()
157157
defer c.mu.Unlock()
158-
if c.patchWorkspaceLogs != nil {
159-
return c.patchWorkspaceLogs()
158+
if c.PatchWorkspaceLogs != nil {
159+
return c.PatchWorkspaceLogs()
160160
}
161161
c.logs = append(c.logs, logs.Logs...)
162162
return nil

‎coderd/tailnet.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ func (s *ServerTailnet) DialAgentNetConn(ctx context.Context, agentID uuid.UUID,
303303
}
304304

305305
node, err := s.getNode(agentID)
306-
if !reachable {
306+
if err != nil {
307307
return nil, xerrors.New("get agent node")
308308
}
309309

0 commit comments

Comments
 (0)