Merge pull request #91 from winebarrel/fix_tty_check

Check stderr tty instead of stdin tty

CHANGELOG.md

@@ -4,6 +4,12 @@
 
 * N/A
 
+## [1.0.3] - 2023-11-27
+
+### Changed
+
+- Check stdout tty instead of stdin tty.
+
 ## [1.0.2] - 2023-11-27
 
 ### Changed

options.go

@@ -8,7 +8,6 @@ import (
 	"github.com/go-sql-driver/mysql"
 	"github.com/jackc/pgx/v5"
 	"github.com/mattn/go-isatty"
-	"github.com/winebarrel/qube/util"
 )
 
 type Options struct {
@@ -27,7 +26,7 @@ type Options struct {
 func (options *Options) AfterApply() error {
 	options.X_Time = JSONDuration(options.Time)
 	options.NullDBOut = os.Stderr
-	options.Progress = isatty.IsTerminal(util.Stdin)
+	options.Progress = isatty.IsTerminal(os.Stderr.Fd())
 
 	if _, err := mysql.ParseDSN(options.DSN); err == nil {
 		options.Driver = DBDriverMySQL

progress.go

@@ -15,16 +15,21 @@ const (
 	InterimReportIntvl = 1 * time.Second
 )
 
+type TTY interface {
+	io.Writer
+	Fd() uintptr
+}
+
 type Progress struct {
-	w           io.Writer
+	w           TTY
 	noop        bool
 	prevDPLen   int
 	nDeadAgents atomic.Uint64
 	closed      chan struct{}
 	startedAt   atomic.Pointer[time.Time] // Use atomic to avoid race conditions
 }
 
-func NewProgress(w io.Writer, noop bool) *Progress {
+func NewProgress(w TTY, noop bool) *Progress {
 	progress := &Progress{
 		w:    w,
 		noop: noop,
@@ -93,7 +98,7 @@ func (progress *Progress) report(rec *Recorder) {
 
 	running := rec.Nagents - progress.nDeadAgents.Load()
 	line := fmt.Sprintf("%02d:%02d | %d agents / exec %d queries, %d errors (%.0f qps)", minute, second, running, rec.CountAll(), rec.ErrorQueryCount(), qps)
-	fmt.Fprintf(progress.w, "\r%-*s", util.MustGetTermSize(), line)
+	fmt.Fprintf(progress.w, "\r%-*s", util.MustGetTermSize(progress.w.Fd()), line)
 }
 
 func (progress *Progress) Close() {
@@ -102,5 +107,5 @@ func (progress *Progress) Close() {
 	}
 
 	<-progress.closed
-	fmt.Fprintf(progress.w, "\r"+strings.Repeat(" ", util.MustGetTermSize())+"\r")
+	fmt.Fprintf(progress.w, "\r"+strings.Repeat(" ", util.MustGetTermSize(progress.w.Fd()))+"\r")
 }

progress_test.go

@@ -10,21 +10,16 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/winebarrel/qube"
-	"github.com/winebarrel/qube/util"
 )
 
 func Test_Progress(t *testing.T) {
-	origStdin := util.Stdin
-	defer func() { util.Stdin = origStdin }()
-
 	require := require.New(t)
 	assert := assert.New(t)
 
 	pty, tty, err := pty.Open()
 	require.NoError(err)
 	defer pty.Close()
 	defer tty.Close()
-	util.Stdin = tty.Fd()
 
 	rec := qube.NewRecorder(testUUID, &qube.Options{})
 	progress := qube.NewProgress(tty, false)

util/term.go

@@ -1,17 +1,11 @@
 package util
 
 import (
-	"os"
-
 	"golang.org/x/term"
 )
 
-var (
-	Stdin = os.Stdin.Fd()
-)
-
-func MustGetTermSize() int {
-	width, _, err := term.GetSize(int(Stdin))
+func MustGetTermSize(fd uintptr) int {
+	width, _, err := term.GetSize(int(fd))
 
 	if err != nil {
 		panic(err)