Add motd and hushlogin test

mafredri · mafredri · commit 61613a141275 · 2022-11-22T15:41:13.000Z
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2,6 +2,7 @@ package agent_test
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -193,6 +194,86 @@ func TestAgent(t *testing.T) {
 		}
 	})
 
+	t.Run("Session TTY MOTD", func(t *testing.T) {
+		t.Parallel()
+
+		if runtime.GOOS == "windows" {
+			// This might be our implementation, or ConPTY itself.
+			// It's difficult to find extensive tests for it, so
+			// it seems like it could be either.
+			t.Skip("ConPTY appears to be inconsistent on Windows.")
+		}
+
+		wantMOTD := "Welcome to your Coder workspace!"
+
+		name := filepath.Join(t.TempDir(), "motd")
+		err := os.WriteFile(name, []byte(wantMOTD), 0o600)
+		require.NoError(t, err, "write motd file")
+
+		session := setupSSHSession(t, codersdk.WorkspaceAgentMetadata{
+			MOTDFile: name,
+		})
+		err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
+		require.NoError(t, err)
+
+		ptty := ptytest.New(t)
+		var stdout bytes.Buffer
+		session.Stdout = &stdout
+		session.Stderr = ptty.Output()
+		session.Stdin = ptty.Input()
+		err = session.Start("")
+		require.NoError(t, err)
+
+		ptty.WriteLine("exit")
+		err = session.Wait()
+		require.NoError(t, err)
+
+		require.Contains(t, stdout.String(), wantMOTD, "should show motd")
+	})
+
+	//nolint:paralleltest // This test sets an environment variable.
+	t.Run("Session TTY Hushlogin", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			// This might be our implementation, or ConPTY itself.
+			// It's difficult to find extensive tests for it, so
+			// it seems like it could be either.
+			t.Skip("ConPTY appears to be inconsistent on Windows.")
+		}
+
+		wantNotMOTD := "Welcome to your Coder workspace!"
+
+		tmpdir := t.TempDir()
+		name := filepath.Join(tmpdir, "motd")
+		err := os.WriteFile(name, []byte(wantNotMOTD), 0o600)
+		require.NoError(t, err, "write motd file")
+		f, err := os.Create(filepath.Join(tmpdir, ".hushlogin"))
+		require.NoError(t, err, "create .hushlogin file")
+		err = f.Close()
+		require.NoError(t, err, "close .hushlogin file")
+
+		t.Setenv("HOME", tmpdir)
+
+		session := setupSSHSession(t, codersdk.WorkspaceAgentMetadata{
+			MOTDFile: name,
+		})
+		err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
+		require.NoError(t, err)
+
+		ptty := ptytest.New(t)
+		var stdout bytes.Buffer
+		session.Stdout = &stdout
+		session.Stderr = ptty.Output()
+		session.Stdin = ptty.Input()
+		err = session.Start("")
+		require.NoError(t, err)
+
+		ptty.WriteLine("exit")
+		err = session.Wait()
+		require.NoError(t, err)
+
+		require.NotContains(t, stdout.String(), wantNotMOTD, "should not show motd")
+	})
+
 	t.Run("LocalForwarding", func(t *testing.T) {
 		t.Parallel()
 		random, err := net.Listen("tcp", "127.0.0.1:0")
