« Back to Index

Go: mocking multiple stdin prompts

View original Gist on GitHub

Tags: #go #tty

1. mocking multiple stdin prompts.go

// https://play.golang.org/p/w0VPpYMytnG
//
// NOTE: This is a correct way to implement, but a quicker/smaller solution would have been not to create the `bufio.NewScanner` instance (unfortunately the `input()` implementation was legacy code that couldn't be updated).

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"
	"time"
)

func input(r io.Reader) (string, error) {
	s := bufio.NewScanner(r)
	for {
		if ok := s.Scan(); !ok {
			return "", s.Err()
		}
		line := strings.TrimSpace(s.Text())
		return line, nil
	}
}

func run(r io.Reader, w io.Writer) error {
	s, err := input(r)
	if err != nil {
		return err
	}
	fmt.Printf("first input(): %+v\n", s)
	fmt.Fprintf(w, "%+v\n", s)

	s, err = input(r)
	if err != nil {
		return err
	}
	fmt.Printf("second input(): %+v\n", s)
	fmt.Fprintf(w, "%+v\n", s)

	return nil
}

func main() {
	var stdout bytes.Buffer

	stdin, prompt := io.Pipe()

	// Wait for user input
	inputc := make(chan string)
	go func() {
		for input := range inputc {
			fmt.Fprintln(prompt, input)
		}
	}()

	// Wait for `run()` to read user input
	done := make(chan bool)

	// Call `run()` and wait for response
	go func() {
		err := run(stdin, &stdout)
		if err != nil {
			log.Fatal(err)
		}
		done <- true
	}()

	// User provides input
	//
	// NOTE: Must provide as much input as is expected to be waited on by `run()`.
	//       For example, if `run()` calls `input()` twice, then provide two messages.
	//       Otherwise the select statement will trigger the timeout error.
	inputc <- "foo!"
	inputc <- "bar!"

	// Wait for result
	select {
	case <-done:
		fmt.Print("\ngot result!\n")
	case <-time.After(time.Second):
		fmt.Print("oh no! no result! timeout!\n")
	}

	// Inspect stdout
	fmt.Printf("\nstdout:\n%+v\n", stdout.String())
}

2. alternatives that ignore the constraint of input function.go

// https://play.golang.org/p/fO6yo1r77ym

package main

import (
	"bufio"
	"fmt"
	"io"
	"strings"
)

func input(r io.Reader) (string, error) {
	s := bufio.NewScanner(r)
	for {
		if ok := s.Scan(); !ok {
			return "", s.Err()
		}
		line := strings.TrimSpace(s.Text())
		return line, nil
	}
}

type inputWrapper struct {
	*bufio.Scanner
	reused *strings.Reader
}

func (i inputWrapper) More() (string, error) {
	if !i.Scanner.Scan() {
		return "", i.Scanner.Err()
	}
	line := i.Scanner.Text()
	if i.reused == nil {
		i.reused = strings.NewReader(line)
	}
	i.reused.Reset(line)
	return input(i.reused)
}

func main() {

	r := strings.NewReader("example.com\ngoogle.com")
	input := inputWrapper{Scanner: bufio.NewScanner(r)}.More

	fmt.Println(input())
	fmt.Println(input())
}