-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathagent_test.go
More file actions
110 lines (103 loc) · 2.82 KB
/
agent_test.go
File metadata and controls
110 lines (103 loc) · 2.82 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package agent_test
import (
"context"
"runtime"
"strings"
"testing"
"github.com/pion/webrtc/v3"
"github.com/stretchr/testify/require"
"go.uber.org/goleak"
"golang.org/x/crypto/ssh"
"cdr.dev/slog"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/agent"
"github.com/coder/coder/peer"
"github.com/coder/coder/peerbroker"
"github.com/coder/coder/peerbroker/proto"
"github.com/coder/coder/provisionersdk"
"github.com/coder/coder/pty/ptytest"
)
func TestMain(m *testing.M) {
goleak.VerifyTestMain(m)
}
func TestAgent(t *testing.T) {
t.Parallel()
t.Run("SessionExec", func(t *testing.T) {
t.Parallel()
api := setup(t)
stream, err := api.NegotiateConnection(context.Background())
require.NoError(t, err)
conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{}, &peer.ConnOptions{
Logger: slogtest.Make(t, nil),
})
require.NoError(t, err)
t.Cleanup(func() {
_ = conn.Close()
})
sshClient, err := agent.DialSSHClient(conn)
require.NoError(t, err)
session, err := sshClient.NewSession()
require.NoError(t, err)
command := "echo test"
if runtime.GOOS == "windows" {
command = "cmd.exe /c echo test"
}
output, err := session.Output(command)
require.NoError(t, err)
require.Equal(t, "test", strings.TrimSpace(string(output)))
})
t.Run("SessionTTY", func(t *testing.T) {
t.Parallel()
api := setup(t)
stream, err := api.NegotiateConnection(context.Background())
require.NoError(t, err)
conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{}, &peer.ConnOptions{
Logger: slogtest.Make(t, nil),
})
require.NoError(t, err)
t.Cleanup(func() {
_ = conn.Close()
})
sshClient, err := agent.DialSSHClient(conn)
require.NoError(t, err)
session, err := sshClient.NewSession()
require.NoError(t, err)
prompt := "$"
command := "bash"
if runtime.GOOS == "windows" {
command = "cmd.exe"
prompt = ">"
}
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
require.NoError(t, err)
ptty := ptytest.New(t)
require.NoError(t, err)
session.Stdout = ptty.Output()
session.Stderr = ptty.Output()
session.Stdin = ptty.Input()
err = session.Start(command)
require.NoError(t, err)
ptty.ExpectMatch(prompt)
ptty.WriteLine("echo test")
ptty.ExpectMatch("test")
ptty.WriteLine("exit")
err = session.Wait()
require.NoError(t, err)
})
}
func setup(t *testing.T) proto.DRPCPeerBrokerClient {
client, server := provisionersdk.TransportPipe()
closer := agent.New(func(ctx context.Context) (*peerbroker.Listener, error) {
return peerbroker.Listen(server, nil, &peer.ConnOptions{
Logger: slogtest.Make(t, nil),
})
}, &agent.Options{
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
})
t.Cleanup(func() {
_ = client.Close()
_ = server.Close()
_ = closer.Close()
})
return proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(client))
}