coder/cli/cliui/agent_test.go

package cliui_test

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"

	"github.com/coder/coder/cli/clibase"
	"github.com/coder/coder/cli/cliui"
	"github.com/coder/coder/codersdk"
	"github.com/coder/coder/pty/ptytest"
	"github.com/coder/coder/testutil"
)

func TestAgent(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	var disconnected atomic.Bool
	ptty := ptytest.New(t)
	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:           codersdk.WorkspaceAgentDisconnected,
						LoginBeforeReady: true,
					}
					if disconnected.Load() {
						agent.Status = codersdk.WorkspaceAgentConnected
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  10 * time.Millisecond,
			})
			return err
		},
	}

	inv := cmd.Invoke()
	ptty.Attach(inv)
	done := make(chan struct{})
	go func() {
		defer close(done)
		err := inv.Run()
		assert.NoError(t, err)
	}()
	ptty.ExpectMatchContext(ctx, "lost connection")
	disconnected.Store(true)
	<-done
}

func TestAgent_TimeoutWithTroubleshootingURL(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	wantURL := "https://coder.com/troubleshoot"

	var connected, timeout atomic.Bool
	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:             codersdk.WorkspaceAgentConnecting,
						TroubleshootingURL: wantURL,
						LoginBeforeReady:   true,
					}
					switch {
					case !connected.Load() && timeout.Load():
						agent.Status = codersdk.WorkspaceAgentTimeout
					case connected.Load():
						agent.Status = codersdk.WorkspaceAgentConnected
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  5 * time.Millisecond,
			})
			return err
		},
	}
	ptty := ptytest.New(t)

	inv := cmd.Invoke()
	ptty.Attach(inv)
	done := make(chan error, 1)
	go func() {
		done <- inv.WithContext(ctx).Run()
	}()
	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
	timeout.Store(true)
	ptty.ExpectMatchContext(ctx, wantURL)
	connected.Store(true)
	require.NoError(t, <-done)
}

func TestAgent_StartupTimeout(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"

	var status, state atomic.String
	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }

	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:             codersdk.WorkspaceAgentConnecting,
						LoginBeforeReady:   false,
						LifecycleState:     codersdk.WorkspaceAgentLifecycleCreated,
						TroubleshootingURL: wantURL,
					}

					if s := status.Load(); s != "" {
						agent.Status = codersdk.WorkspaceAgentStatus(s)
					}
					if s := state.Load(); s != "" {
						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  time.Millisecond,
				NoWait:        false,
			})
			return err
		},
	}

	ptty := ptytest.New(t)

	inv := cmd.Invoke()
	ptty.Attach(inv)
	done := make(chan error, 1)
	go func() {
		done <- inv.WithContext(ctx).Run()
	}()
	setStatus(codersdk.WorkspaceAgentConnecting)
	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
	setStatus(codersdk.WorkspaceAgentConnected)
	setState(codersdk.WorkspaceAgentLifecycleStarting)
	ptty.ExpectMatchContext(ctx, "workspace is getting ready")
	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
	ptty.ExpectMatchContext(ctx, "is taking longer")
	ptty.ExpectMatchContext(ctx, wantURL)
	setState(codersdk.WorkspaceAgentLifecycleReady)
	require.NoError(t, <-done)
}

func TestAgent_StartErrorExit(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"

	var status, state atomic.String
	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:             codersdk.WorkspaceAgentConnecting,
						LoginBeforeReady:   false,
						LifecycleState:     codersdk.WorkspaceAgentLifecycleCreated,
						TroubleshootingURL: wantURL,
					}

					if s := status.Load(); s != "" {
						agent.Status = codersdk.WorkspaceAgentStatus(s)
					}
					if s := state.Load(); s != "" {
						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  60 * time.Second,
				NoWait:        false,
			})
			return err
		},
	}

	ptty := ptytest.New(t)

	inv := cmd.Invoke()
	ptty.Attach(inv)
	done := make(chan error, 1)
	go func() {
		done <- inv.WithContext(ctx).Run()
	}()
	setStatus(codersdk.WorkspaceAgentConnected)
	setState(codersdk.WorkspaceAgentLifecycleStarting)
	ptty.ExpectMatchContext(ctx, "to become ready...")
	setState(codersdk.WorkspaceAgentLifecycleStartError)
	ptty.ExpectMatchContext(ctx, "ran into a problem")
	err := <-done
	require.ErrorIs(t, err, cliui.AgentStartError, "lifecycle start_error should exit with error")
}

func TestAgent_NoWait(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"

	var status, state atomic.String
	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:             codersdk.WorkspaceAgentConnecting,
						LoginBeforeReady:   false,
						LifecycleState:     codersdk.WorkspaceAgentLifecycleCreated,
						TroubleshootingURL: wantURL,
					}

					if s := status.Load(); s != "" {
						agent.Status = codersdk.WorkspaceAgentStatus(s)
					}
					if s := state.Load(); s != "" {
						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  time.Second,
				NoWait:        true,
			})
			return err
		},
	}

	ptty := ptytest.New(t)

	inv := cmd.Invoke()
	ptty.Attach(inv)
	done := make(chan error, 1)
	go func() {
		done <- inv.WithContext(ctx).Run()
	}()
	setStatus(codersdk.WorkspaceAgentConnecting)
	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")

	setStatus(codersdk.WorkspaceAgentConnected)
	require.NoError(t, <-done, "created - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStarting)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "starting - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "start timeout - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStartError)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "start error - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleReady)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "ready - should exit early")
}

func TestAgent_LoginBeforeReadyEnabled(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
	defer cancel()

	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"

	var status, state atomic.String
	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
	cmd := &clibase.Cmd{
		Handler: func(inv *clibase.Invocation) error {
			err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
				WorkspaceName: "example",
				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
					agent := codersdk.WorkspaceAgent{
						Status:             codersdk.WorkspaceAgentConnecting,
						LoginBeforeReady:   true,
						LifecycleState:     codersdk.WorkspaceAgentLifecycleCreated,
						TroubleshootingURL: wantURL,
					}

					if s := status.Load(); s != "" {
						agent.Status = codersdk.WorkspaceAgentStatus(s)
					}
					if s := state.Load(); s != "" {
						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
					}
					return agent, nil
				},
				FetchInterval: time.Millisecond,
				WarnInterval:  time.Second,
				NoWait:        false,
			})
			return err
		},
	}

	inv := cmd.Invoke()

	ptty := ptytest.New(t)
	ptty.Attach(inv)
	done := make(chan error, 1)
	go func() {
		done <- inv.WithContext(ctx).Run()
	}()
	setStatus(codersdk.WorkspaceAgentConnecting)
	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")

	setStatus(codersdk.WorkspaceAgentConnected)
	require.NoError(t, <-done, "created - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStarting)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "starting - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "start timeout - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleStartError)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "start error - should exit early")

	setState(codersdk.WorkspaceAgentLifecycleReady)
	go func() { done <- inv.WithContext(ctx).Run() }()
	require.NoError(t, <-done, "ready - should exit early")
}