use workspaceagentconn methods added in #12593

coder · johnstcn · Mar 15, 2024 · Mar 12, 2024 · Mar 12, 2024 · Mar 13, 2024
commit df184a09eafb6c89cd5ecb04b66ffbe75a0259a6
diff --git a/cli/support_test.go b/cli/support_test.go
@@ -4,6 +4,7 @@ import (
 	"archive/zip"
 	"encoding/json"
 	"io"
+	"os"
 	"path/filepath"
 	"runtime"
 	"testing"
@@ -13,6 +14,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/cli/clitest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -42,7 +44,12 @@ func TestSupportBundle(t *testing.T) {
 		}).WithAgent().Do()
 		ws, err := client.Workspace(ctx, r.Workspace.ID)
 		require.NoError(t, err)
-		agt := agenttest.New(t, client.URL, r.AgentToken)
+		tempDir := t.TempDir()
+		logPath := filepath.Join(tempDir, "coder-agent.log")
+		require.NoError(t, os.WriteFile(logPath, []byte("hello from the agent"), 0o600))
+		agt := agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
+			o.LogDir = tempDir
+		})
 		defer agt.Close()
 		coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).Wait()
 

diff --git a/support/support.go b/support/support.go
@@ -356,29 +356,21 @@ func AgentInfo(ctx context.Context, client *codersdk.Client, log slog.Logger, ag
 			})
 
 			eg.Go(func() error {
-				tokenBytes, err := runCmd(ctx, client, agentID, `echo $CODER_AGENT_TOKEN`)
-				if err != nil {
-					return xerrors.Errorf("failed to fetch agent token: %w", err)
-				}
-				token := string(bytes.TrimSpace(tokenBytes))
-				agentClient := agentsdk.New(client.URL)
-				agentClient.SetSessionToken(token)
-				manifestRes, err := agentClient.SDK.Request(ctx, http.MethodGet, "/api/v2/workspaceagents/me/manifest", nil)
+				manifestRes, err := conn.DebugManifest(ctx)
 				if err != nil {
 					return xerrors.Errorf("fetch manifest: %w", err)
 				}
-				defer manifestRes.Body.Close()
-				if err := json.NewDecoder(manifestRes.Body).Decode(&a.Manifest); err != nil {
+				if err := json.NewDecoder(bytes.NewReader(manifestRes)).Decode(&a.Manifest); err != nil {
 					return xerrors.Errorf("decode agent manifest: %w", err)
 				}
 
 				return nil
 			})
 
 			eg.Go(func() error {
-				logBytes, err := runCmd(ctx, client, agentID, `tail -n 1000 /tmp/coder-agent.log`)
+				logBytes, err := conn.DebugLogs(ctx)
 				if err != nil {
-					return xerrors.Errorf("tail /tmp/coder-agent.log: %w", err)
+					return xerrors.Errorf("fetch coder agent logs: %w", err)
 				}
 				a.Logs = logBytes
 				return nil
@@ -472,30 +464,3 @@ func sanitizeEnv(kvs map[string]string) {
 		}
 	}
 }
-
-// TODO: use rpty instead? which is less liable to fail?
-func runCmd(ctx context.Context, client *codersdk.Client, agentID uuid.UUID, cmd string) ([]byte, error) {
-	agentConn, err := client.DialWorkspaceAgent(ctx, agentID, &codersdk.DialWorkspaceAgentOptions{})
-	if err != nil {
-		return nil, xerrors.Errorf("dial workspace agent: %w", err)
-	}
-	defer agentConn.Close()
-
-	sshClient, err := agentConn.SSHClient(ctx)
-	if err != nil {
-		return nil, xerrors.Errorf("get ssh client: %w", err)
-	}
-	defer sshClient.Close()
-
-	sshSession, err := sshClient.NewSession()
-	if err != nil {
-		return nil, xerrors.Errorf("new ssh session: %w", err)
-	}
-	defer sshSession.Close()
-
-	cmdBytes, err := sshSession.CombinedOutput(cmd)
-	if err != nil {
-		return nil, xerrors.Errorf("shell: %w", err)
-	}
-	return cmdBytes, err
-}
diff --git a/support/support_test.go b/support/support_test.go
@@ -5,6 +5,8 @@ import (
 	"context"
 	"io"
 	"net/http"
+	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
@@ -15,6 +17,7 @@ import (
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
@@ -212,7 +215,12 @@ func setupWorkspaceAndAgent(ctx context.Context, t *testing.T, client *codersdk.
 	})
 	require.NoError(t, err)
 
-	_ = agenttest.New(t, client.URL, wbr.AgentToken)
+	tempDir := t.TempDir()
+	logPath := filepath.Join(tempDir, "coder-agent.log")
+	require.NoError(t, os.WriteFile(logPath, []byte("hello from the agent"), 0o600))
+	_ = agenttest.New(t, client.URL, wbr.AgentToken, func(o *agent.Options) {
+		o.LogDir = tempDir
+	})
 	coderdtest.NewWorkspaceAgentWaiter(t, client, wbr.Workspace.ID).Wait()
 
 	return ws, agt