mirror of https://github.com/coder/coder.git
142 lines
3.1 KiB
Go
142 lines
3.1 KiB
Go
package pty_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/coder/coder/v2/pty"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
// Test_Start_copy tests that we can use io.Copy() on command output
|
|
// without deadlocking.
|
|
func Test_Start_copy(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
|
defer cancel()
|
|
|
|
pc, cmd, err := pty.Start(pty.CommandContext(ctx, cmdEcho, argEcho...))
|
|
require.NoError(t, err)
|
|
b := &bytes.Buffer{}
|
|
readDone := make(chan error, 1)
|
|
go func() {
|
|
_, err := io.Copy(b, pc.OutputReader())
|
|
readDone <- err
|
|
}()
|
|
|
|
select {
|
|
case err := <-readDone:
|
|
require.NoError(t, err)
|
|
case <-ctx.Done():
|
|
t.Error("read timed out")
|
|
}
|
|
assert.Contains(t, b.String(), "test")
|
|
|
|
cmdDone := make(chan error, 1)
|
|
go func() {
|
|
cmdDone <- cmd.Wait()
|
|
}()
|
|
|
|
select {
|
|
case err := <-cmdDone:
|
|
require.NoError(t, err)
|
|
case <-ctx.Done():
|
|
t.Error("cmd.Wait() timed out")
|
|
}
|
|
}
|
|
|
|
// Test_Start_truncation tests that we can read command output without truncation
|
|
// even after the command has exited.
|
|
func Test_Start_truncation(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
|
|
defer cancel()
|
|
|
|
pc, cmd, err := pty.Start(pty.CommandContext(ctx, cmdCount, argCount...))
|
|
|
|
require.NoError(t, err)
|
|
readDone := make(chan struct{})
|
|
go func() {
|
|
defer close(readDone)
|
|
terminalReader := testutil.NewTerminalReader(t, pc.OutputReader())
|
|
n := 1
|
|
for n <= countEnd {
|
|
want := fmt.Sprintf("%d", n)
|
|
err := terminalReader.ReadUntilString(ctx, want)
|
|
assert.NoError(t, err, "want: %s", want)
|
|
if err != nil {
|
|
return
|
|
}
|
|
n++
|
|
if (countEnd - n) < 100 {
|
|
// If the OS buffers the output, the process can exit even if
|
|
// we're not done reading. We want to slow our reads so that
|
|
// if there is a race between reading the data and it being
|
|
// truncated, we will lose and fail the test.
|
|
time.Sleep(testutil.IntervalFast)
|
|
}
|
|
}
|
|
// ensure we still get to EOF
|
|
endB := &bytes.Buffer{}
|
|
_, err := io.Copy(endB, pc.OutputReader())
|
|
assert.NoError(t, err)
|
|
}()
|
|
|
|
cmdDone := make(chan error, 1)
|
|
go func() {
|
|
cmdDone <- cmd.Wait()
|
|
}()
|
|
|
|
select {
|
|
case err := <-cmdDone:
|
|
require.NoError(t, err)
|
|
case <-ctx.Done():
|
|
t.Fatal("cmd.Wait() timed out")
|
|
}
|
|
|
|
select {
|
|
case <-readDone:
|
|
// OK!
|
|
case <-ctx.Done():
|
|
t.Fatal("read timed out")
|
|
}
|
|
}
|
|
|
|
// Test_Start_cancel_context tests that we can cancel the command context and kill the process.
|
|
func Test_Start_cancel_context(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
|
|
defer cancel()
|
|
cmdCtx, cmdCancel := context.WithCancel(ctx)
|
|
|
|
pc, cmd, err := pty.Start(pty.CommandContext(cmdCtx, cmdSleep, argSleep...))
|
|
require.NoError(t, err)
|
|
defer func() {
|
|
_ = pc.Close()
|
|
}()
|
|
cmdCancel()
|
|
|
|
cmdDone := make(chan struct{})
|
|
go func() {
|
|
defer close(cmdDone)
|
|
_ = cmd.Wait()
|
|
}()
|
|
|
|
select {
|
|
case <-cmdDone:
|
|
// OK!
|
|
case <-ctx.Done():
|
|
t.Error("cmd.Wait() timed out")
|
|
}
|
|
}
|