From 3b0b650318936582f76dc95924ef2a239e27edfe Mon Sep 17 00:00:00 2001
From: Ryo Ota <nwtgck@nwtgck.org>
Date: Fri, 11 Aug 2023 00:30:26 +0900
Subject: [PATCH] test pty

---
 cmd/root_test.go      | 12 ++++++------
 cmd/test_util_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/cmd/root_test.go b/cmd/root_test.go
index 1f292c9..dc00f86 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -57,10 +57,10 @@ func TestAllPermissionsAllowed(t *testing.T) {
 	assert.NoError(t, err)
 	defer client.Close()
 	assert.NoError(t, err)
-	assertExec(t, client)
-	assertLocalPortForwarding(t, client)
 	assertRemotePortForwardingTODO(t, client)
-	// TODO: pty
+	assertLocalPortForwarding(t, client)
+	assertExec(t, client)
+	assertPtyTerminal(t, client)
 	// TODO: sftp
 }
 
@@ -172,7 +172,7 @@ func TestAllowExecute(t *testing.T) {
 	assertNoRemotePortForwarding(t, client)
 	assertNoLocalPortForwarding(t, client)
 	assertExec(t, client)
-	// TODO: no pty
+	assertPtyTerminal(t, client)
 	// TODO: no sftp
 }
 
@@ -203,7 +203,7 @@ func TestAllowTcpipForward(t *testing.T) {
 	assertRemotePortForwardingTODO(t, client)
 	assertNoLocalPortForwarding(t, client)
 	assertNoExec(t, client)
-	// TODO: no pty
+	assertNoPtyTerminal(t, client)
 	// TODO: no sftp
 }
 
@@ -234,7 +234,7 @@ func TestAllowDirectTcpip(t *testing.T) {
 	assertNoRemotePortForwarding(t, client)
 	assertLocalPortForwarding(t, client)
 	assertNoExec(t, client)
-	// TODO: no pty
+	assertNoPtyTerminal(t, client)
 	// TODO: no sftp
 }
 
diff --git a/cmd/test_util_test.go b/cmd/test_util_test.go
index 53b3450..016a7f5 100644
--- a/cmd/test_util_test.go
+++ b/cmd/test_util_test.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"bytes"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/crypto/ssh"
 	"io"
@@ -8,6 +9,7 @@ import (
 	"os/exec"
 	"strconv"
 	"testing"
+	"time"
 )
 
 func getAvailableTcpPort() int {
@@ -49,6 +51,43 @@ func assertNoExec(t *testing.T, client *ssh.Client) {
 	assert.Equal(t, "ssh: command whoami failed", err.Error())
 }
 
+func assertPtyTerminal(t *testing.T, client *ssh.Client) {
+	session, err := client.NewSession()
+	assert.NoError(t, err)
+	defer session.Close()
+
+	err = session.RequestPty("xterm", 100, 200, ssh.TerminalModes{})
+	assert.NoError(t, err)
+	stdin, err := session.StdinPipe()
+	assert.NoError(t, err)
+	_, err = stdin.Write([]byte("echo helloworldviapty\r"))
+	assert.NoError(t, err)
+	stdout, err := session.StdoutPipe()
+	assert.NoError(t, err)
+	stdoutBytesChan := make(chan []byte)
+	go func() {
+		var buff bytes.Buffer
+		_, err := io.Copy(&buff, stdout)
+		assert.NoError(t, err)
+		stdoutBytesChan <- buff.Bytes()
+	}()
+	err = session.Shell()
+	assert.NoError(t, err)
+	time.Sleep(1 * time.Second)
+	session.Close()
+	stdoutBytes := <-stdoutBytesChan
+	assert.Contains(t, string(stdoutBytes), "helloworldviapty")
+}
+
+func assertNoPtyTerminal(t *testing.T, client *ssh.Client) {
+	session, err := client.NewSession()
+	assert.NoError(t, err)
+	defer session.Close()
+	err = session.RequestPty("xterm", 100, 200, ssh.TerminalModes{})
+	assert.Error(t, err)
+	assert.Equal(t, "ssh: pty-req failed", err.Error())
+}
+
 func assertLocalPortForwarding(t *testing.T, client *ssh.Client) {
 	var remoteTcpPort int
 	acceptedConnChan := make(chan net.Conn)