Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 5d9b4d2

Browse files
coadlercoadler
committed
fix: log after test exit in TestAgent/StartupScript 
1 parent 104c76b commit 5d9b4d2

File tree

4 files changed

+35
-17
lines changed

4 files changed

+35
-17
lines changed

‎agent/agent.go

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,22 +21,19 @@ import (
2121
"time"
2222

2323
"github.com/armon/circbuf"
24+
"github.com/gliderlabs/ssh"
2425
"github.com/google/uuid"
25-
26+
"github.com/pkg/sftp"
2627
"go.uber.org/atomic"
28+
gossh "golang.org/x/crypto/ssh"
29+
"golang.org/x/xerrors"
2730

2831
"cdr.dev/slog"
2932
"github.com/coder/coder/agent/usershell"
3033
"github.com/coder/coder/peer"
3134
"github.com/coder/coder/peerbroker"
3235
"github.com/coder/coder/pty"
3336
"github.com/coder/retry"
34-
35-
"github.com/pkg/sftp"
36-
37-
"github.com/gliderlabs/ssh"
38-
gossh "golang.org/x/crypto/ssh"
39-
"golang.org/x/xerrors"
4037
)
4138

4239
type Options struct {
@@ -174,17 +171,25 @@ func (*agent) runStartupScript(ctx context.Context, script string) error {
174171
defer func() {
175172
_ = writer.Close()
176173
}()
174+
177175
caller := "-c"
178176
if runtime.GOOS == "windows" {
179177
caller = "/c"
180178
}
179+
181180
cmd := exec.CommandContext(ctx, shell, caller, script)
182181
cmd.Stdout = writer
183182
cmd.Stderr = writer
184183
err = cmd.Run()
185184
if err != nil {
185+
// cmd.Run does not return a context canceled error, it returns "signal: killed".
186+
if ctx.Err() != nil {
187+
return ctx.Err()
188+
}
189+
186190
return xerrors.Errorf("run: %w", err)
187191
}
192+
188193
return nil
189194
}
190195

@@ -208,11 +213,11 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
208213
}
209214

210215
switch channel.Protocol() {
211-
case "ssh":
216+
case peer.ProtocolSSH:
212217
go a.sshServer.HandleConn(channel.NetConn())
213-
case "reconnecting-pty":
218+
case peer.ProtocolReconnectingPTY:
214219
go a.handleReconnectingPTY(ctx, channel.Label(), channel.NetConn())
215-
case "dial":
220+
case peer.ProtocolDial:
216221
go a.handleDial(ctx, channel.Label(), channel.NetConn())
217222
default:
218223
a.logger.Warn(ctx, "unhandled protocol from channel",
@@ -478,8 +483,8 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
478483
a.logger.Warn(ctx, "start reconnecting pty command", slog.F("id", id))
479484
}
480485

481-
// Default to buffer 64KB.
482-
circularBuffer, err := circbuf.NewBuffer(64 * 1024)
486+
// Default to buffer 64KiB.
487+
circularBuffer, err := circbuf.NewBuffer(64 << 10)
483488
if err != nil {
484489
a.logger.Warn(ctx, "create circular buffer", slog.Error(err))
485490
return

‎agent/agent_test.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -172,8 +172,9 @@ func TestAgent(t *testing.T) {
172172
tempPath := filepath.Join(os.TempDir(), "content.txt")
173173
content := "somethingnice"
174174
setupAgent(t, agent.Metadata{
175-
StartupScript: "echo " + content + " > " + tempPath,
175+
StartupScript: fmt.Sprintf("echo %s > %s", content, tempPath),
176176
}, 0)
177+
177178
var gotContent string
178179
require.Eventually(t, func() bool {
179180
content, err := os.ReadFile(tempPath)
@@ -202,7 +203,12 @@ func TestAgent(t *testing.T) {
202203
// it seems like it could be either.
203204
t.Skip("ConPTY appears to be inconsistent on Windows.")
204205
}
205-
conn := setupAgent(t, agent.Metadata{}, 0)
206+
207+
conn := setupAgent(t, agent.Metadata{
208+
EnvironmentVariables: map[string]string{
209+
// "SHELL": "/bin/bash",
210+
},
211+
}, 0)
206212
id := uuid.NewString()
207213
netConn, err := conn.ReconnectingPTY(id, 100, 100)
208214
require.NoError(t, err)
@@ -228,6 +234,7 @@ func TestAgent(t *testing.T) {
228234
}
229235
}
230236
}
237+
231238
matchEchoCommand := func(line string) bool {
232239
return strings.Contains(line, "echo test")
233240
}

‎agent/conn.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ type Conn struct {
3636
// be reconnected to via ID.
3737
func (c *Conn) ReconnectingPTY(id string, height, width uint16) (net.Conn, error) {
3838
channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d", id, height, width), &peer.ChannelOptions{
39-
Protocol: "reconnecting-pty",
39+
Protocol: peer.ProtocolReconnectingPTY,
4040
})
4141
if err != nil {
4242
return nil, xerrors.Errorf("pty: %w", err)
@@ -47,7 +47,7 @@ func (c *Conn) ReconnectingPTY(id string, height, width uint16) (net.Conn, error
4747
// SSH dials the built-in SSH server.
4848
func (c *Conn) SSH() (net.Conn, error) {
4949
channel, err := c.CreateChannel(context.Background(), "ssh", &peer.ChannelOptions{
50-
Protocol: "ssh",
50+
Protocol: peer.ProtocolSSH,
5151
})
5252
if err != nil {
5353
return nil, xerrors.Errorf("dial: %w", err)
@@ -87,7 +87,7 @@ func (c *Conn) DialContext(ctx context.Context, network string, addr string) (ne
8787
}
8888

8989
channel, err := c.CreateChannel(ctx, u.String(), &peer.ChannelOptions{
90-
Protocol: "dial",
90+
Protocol: peer.ProtocolDial,
9191
Unordered: strings.HasPrefix(network, "udp"),
9292
})
9393
if err != nil {

‎peer/channel.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ func newChannel(conn *Conn, dc *webrtc.DataChannel, opts *ChannelOptions) *Chann
4141
return channel
4242
}
4343

44+
const (
45+
ProtocolReconnectingPTY = "reconnecting-pty"
46+
ProtocolSSH = "ssh"
47+
ProtocolDial = "dial"
48+
)
49+
4450
type ChannelOptions struct {
4551
// ID is a channel ID that should be used when `Negotiated`
4652
// is true.

0 commit comments

Comments
 (0)