dnote/pkg/assert/prompt.go
2025-10-31 23:41:21 -07:00

84 lines
2.2 KiB
Go

/* Copyright 2025 Dnote Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package assert
import (
"bufio"
"io"
"strings"
"time"
"github.com/pkg/errors"
)
// WaitForPrompt waits for an expected prompt to appear in stdout with a timeout.
// Returns an error if the prompt is not found within the timeout period.
// Handles prompts with or without newlines by reading character by character.
func WaitForPrompt(stdout io.Reader, expectedPrompt string, timeout time.Duration) error {
type result struct {
found bool
err error
}
resultCh := make(chan result, 1)
go func() {
reader := bufio.NewReader(stdout)
var buffer strings.Builder
found := false
for {
b, err := reader.ReadByte()
if err != nil {
resultCh <- result{found: found, err: err}
return
}
buffer.WriteByte(b)
if strings.Contains(buffer.String(), expectedPrompt) {
found = true
break
}
}
resultCh <- result{found: found, err: nil}
}()
select {
case res := <-resultCh:
if res.err != nil && res.err != io.EOF {
return errors.Wrap(res.err, "reading stdout")
}
if !res.found {
return errors.Errorf("expected prompt '%s' not found in stdout", expectedPrompt)
}
return nil
case <-time.After(timeout):
return errors.Errorf("timeout waiting for prompt '%s'", expectedPrompt)
}
}
// RespondToPrompt is a helper that waits for a prompt and sends a response.
func RespondToPrompt(stdout io.Reader, stdin io.WriteCloser, expectedPrompt, response string, timeout time.Duration) error {
if err := WaitForPrompt(stdout, expectedPrompt, timeout); err != nil {
return err
}
if _, err := io.WriteString(stdin, response); err != nil {
return errors.Wrap(err, "writing response to stdin")
}
return nil
}