waveterm/cmd/test-conn/testutil.go
Mike Sawka 16877fe433
Working on Test Harness for Remote Connections (#2829)
This pull request introduces a new `test-conn` command-line tool for
testing SSH connection flows and user input handling, along with
improvements to user input provider extensibility and several
configuration and validation enhancements.

* New `test-conn` CLI Tool
* User Input Provider Abstraction
* Refactored user input handling to use a pluggable `UserInputProvider`
interface, with a default `FrontendProvider` and the ability to set a
custom provider
* AI Configuration and Verbosity updates
* Enforced that a `SwapToken` must be present in `CommandOptsType` when
starting a remote shell with wsh, improving validation and error
handling. (`pkg/shellexec/shellexec.go`)
* Improved config directory watcher logic to log the config directory
path and avoid logging errors for non-existent subdirectories
2026-02-09 21:50:34 -08:00

326 lines
8.4 KiB
Go

// Copyright 2026, Command Line Inc.
// SPDX-License-Identifier: Apache-2.0
package main
import (
"context"
"fmt"
"log"
"os"
"path/filepath"
"runtime"
"time"
"github.com/google/uuid"
"github.com/wavetermdev/waveterm/pkg/remote"
"github.com/wavetermdev/waveterm/pkg/remote/conncontroller"
"github.com/wavetermdev/waveterm/pkg/shellexec"
"github.com/wavetermdev/waveterm/pkg/userinput"
"github.com/wavetermdev/waveterm/pkg/util/shellutil"
"github.com/wavetermdev/waveterm/pkg/wavebase"
"github.com/wavetermdev/waveterm/pkg/wavejwt"
"github.com/wavetermdev/waveterm/pkg/waveobj"
"github.com/wavetermdev/waveterm/pkg/wconfig"
"github.com/wavetermdev/waveterm/pkg/wshrpc/wshserver"
"github.com/wavetermdev/waveterm/pkg/wshutil"
"github.com/wavetermdev/waveterm/pkg/wstore"
)
func setupWaveEnvVars() error {
homeDir, err := os.UserHomeDir()
if err != nil {
return fmt.Errorf("failed to get home directory: %w", err)
}
isDev := os.Getenv("WAVETERM_DEV") != ""
devSuffix := ""
if isDev {
devSuffix = "-dev"
}
configHome := os.Getenv("WAVETERM_CONFIG_HOME")
if configHome == "" {
configHome = filepath.Join(homeDir, ".config", "waveterm"+devSuffix)
os.Setenv("WAVETERM_CONFIG_HOME", configHome)
}
log.Printf("Using config directory: %s", configHome)
dataHome := os.Getenv("WAVETERM_DATA_HOME")
if dataHome == "" {
if runtime.GOOS == "darwin" {
dataHome = filepath.Join(homeDir, "Library", "Application Support", "waveterm"+devSuffix)
os.Setenv("WAVETERM_DATA_HOME", dataHome)
} else {
return fmt.Errorf("WAVETERM_DATA_HOME must be set on non-macOS systems")
}
}
log.Printf("Using data directory: %s", dataHome)
return nil
}
func initTestHarness(autoAccept bool) error {
log.Printf("Initializing test harness...")
err := setupWaveEnvVars()
if err != nil {
return fmt.Errorf("failed to setup wave env vars: %w", err)
}
err = wavebase.CacheAndRemoveEnvVars()
if err != nil {
return fmt.Errorf("failed to cache env vars: %w", err)
}
wshutil.DefaultRouter = wshutil.NewWshRouter()
wshutil.DefaultRouter.SetAsRootRouter()
wstore.SetClientId("test-client-" + fmt.Sprintf("%d", time.Now().Unix()))
userinput.SetUserInputProvider(&CLIProvider{AutoAccept: autoAccept})
keyPair, err := wavejwt.GenerateKeyPair()
if err != nil {
return fmt.Errorf("failed to generate JWT key pair: %w", err)
}
err = wavejwt.SetPrivateKey(keyPair.PrivateKey)
if err != nil {
return fmt.Errorf("failed to set JWT private key: %w", err)
}
err = wavejwt.SetPublicKey(keyPair.PublicKey)
if err != nil {
return fmt.Errorf("failed to set JWT public key: %w", err)
}
rpc := wshserver.GetMainRpcClient()
wshutil.DefaultRouter.RegisterTrustedLeaf(rpc, wshutil.DefaultRoute)
wconfig.GetWatcher().Start()
log.Printf("Test harness initialized")
return nil
}
func testBasicConnect(connName string, timeout time.Duration) error {
opts, err := remote.ParseOpts(connName)
if err != nil {
return fmt.Errorf("failed to parse connection string: %w", err)
}
log.Printf("Connecting to %s...", opts.String())
conn := conncontroller.GetConn(opts)
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
err = conn.Connect(ctx, &wconfig.ConnKeywords{})
if err != nil {
return fmt.Errorf("connection failed: %w", err)
}
status := conn.DeriveConnStatus()
log.Printf("✓ Connected!")
log.Printf(" Status: %s", status.Status)
log.Printf(" WshEnabled: %v", status.WshEnabled)
log.Printf(" Connection: %s", status.Connection)
if status.WshVersion != "" {
log.Printf(" WshVersion: %s", status.WshVersion)
}
if status.WshError != "" {
log.Printf(" WshError: %s", status.WshError)
}
if status.NoWshReason != "" {
log.Printf(" NoWshReason: %s", status.NoWshReason)
}
return nil
}
func testShellWithCommand(connName string, cmd string, timeout time.Duration) error {
opts, err := remote.ParseOpts(connName)
if err != nil {
return fmt.Errorf("failed to parse connection string: %w", err)
}
log.Printf("Connecting to %s...", opts.String())
conn := conncontroller.GetConn(opts)
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
err = conn.Connect(ctx, &wconfig.ConnKeywords{})
if err != nil {
return fmt.Errorf("connection failed: %w", err)
}
log.Printf("✓ Connected! Starting shell...")
termSize := waveobj.TermSize{Rows: 24, Cols: 80}
shellProc, err := shellexec.StartRemoteShellProcNoWsh(ctx, termSize, "", shellexec.CommandOptsType{}, conn)
if err != nil {
return fmt.Errorf("failed to start shell: %w", err)
}
defer shellProc.Close()
log.Printf("✓ Shell started! Executing: %s", cmd)
_, err = shellProc.Cmd.Write([]byte(cmd + "\n"))
if err != nil {
return fmt.Errorf("failed to write command: %w", err)
}
time.Sleep(500 * time.Millisecond)
buf := make([]byte, 8192)
n, err := shellProc.Cmd.Read(buf)
if err != nil {
log.Printf("Warning: read error (may be expected): %v", err)
}
if n > 0 {
log.Printf("\n--- Output ---\n%s\n--- End Output ---", string(buf[:n]))
} else {
log.Printf("No output received (timeout or no data)")
}
return nil
}
func testWshExec(connName string, cmd string, timeout time.Duration) error {
opts, err := remote.ParseOpts(connName)
if err != nil {
return fmt.Errorf("failed to parse connection string: %w", err)
}
log.Printf("Connecting to %s with wsh enabled...", opts.String())
conn := conncontroller.GetConn(opts)
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
wshEnabled := true
err = conn.Connect(ctx, &wconfig.ConnKeywords{
ConnWshEnabled: &wshEnabled,
})
if err != nil {
return fmt.Errorf("connection failed: %w", err)
}
status := conn.DeriveConnStatus()
log.Printf("✓ Connected! (wsh enabled: %v)", status.WshEnabled)
if status.WshVersion != "" {
log.Printf(" wsh version: %s", status.WshVersion)
}
if !status.WshEnabled {
log.Printf(" WARNING: wsh not enabled - reason: %s", status.NoWshReason)
}
log.Printf("Starting wsh-enabled shell...")
swapToken := &shellutil.TokenSwapEntry{
Token: uuid.New().String(),
Env: make(map[string]string),
Exp: time.Now().Add(5 * time.Minute),
}
swapToken.Env["TERM_PROGRAM"] = "waveterm"
swapToken.Env["WAVETERM"] = "1"
swapToken.Env["WAVETERM_VERSION"] = wavebase.WaveVersion
swapToken.Env["WAVETERM_CONN"] = connName
cmdOpts := shellexec.CommandOptsType{
SwapToken: swapToken,
}
termSize := waveobj.TermSize{Rows: 24, Cols: 80}
shellProc, err := shellexec.StartRemoteShellProc(ctx, ctx, termSize, "", cmdOpts, conn)
if err != nil {
return fmt.Errorf("failed to start shell: %w", err)
}
defer shellProc.Close()
log.Printf("✓ Shell started! Executing: %s", cmd)
_, err = shellProc.Cmd.Write([]byte(cmd + "\n"))
if err != nil {
return fmt.Errorf("failed to write command: %w", err)
}
time.Sleep(500 * time.Millisecond)
buf := make([]byte, 8192)
n, err := shellProc.Cmd.Read(buf)
if err != nil {
log.Printf("Warning: read error (may be expected): %v", err)
}
if n > 0 {
log.Printf("\n--- Output ---\n%s\n--- End Output ---", string(buf[:n]))
} else {
log.Printf("No output received (timeout or no data)")
}
return nil
}
func testInteractiveShell(connName string, timeout time.Duration) error {
opts, err := remote.ParseOpts(connName)
if err != nil {
return fmt.Errorf("failed to parse connection string: %w", err)
}
log.Printf("Connecting to %s...", opts.String())
conn := conncontroller.GetConn(opts)
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
err = conn.Connect(ctx, &wconfig.ConnKeywords{})
if err != nil {
return fmt.Errorf("connection failed: %w", err)
}
log.Printf("✓ Connected! Starting interactive shell...")
log.Printf("Note: This is a simple test - output may be mixed with prompts")
log.Printf("Type commands and press Enter. Type 'exit' to quit.\n")
termSize := waveobj.TermSize{Rows: 24, Cols: 80}
shellProc, err := shellexec.StartRemoteShellProcNoWsh(ctx, termSize, "", shellexec.CommandOptsType{}, conn)
if err != nil {
return fmt.Errorf("failed to start shell: %w", err)
}
defer shellProc.Close()
go func() {
buf := make([]byte, 8192)
for {
n, err := shellProc.Cmd.Read(buf)
if err != nil {
return
}
if n > 0 {
fmt.Print(string(buf[:n]))
}
}
}()
go func() {
buf := make([]byte, 1)
for {
n, err := os.Stdin.Read(buf)
if err != nil {
return
}
if n > 0 {
shellProc.Cmd.Write(buf[:n])
}
}
}()
shellProc.Wait()
log.Printf("\nShell exited")
return nil
}