github.com/khulnasoft/cli@v0.0.0-20240402070845-01bcad7beefa/internal/test/cmd.go (about)

     1  package test
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"syscall"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/khulnasoft/cli/cli/streams"
    11  	"github.com/spf13/cobra"
    12  	"gotest.tools/v3/assert"
    13  )
    14  
    15  func TerminatePrompt(ctx context.Context, t *testing.T, cmd *cobra.Command, cli *FakeCli, assertFunc func(*testing.T, error)) {
    16  	t.Helper()
    17  
    18  	errChan := make(chan error)
    19  	defer close(errChan)
    20  
    21  	// wrap the out stream to detect when the prompt is ready
    22  	writerHookChan := make(chan struct{})
    23  	defer close(writerHookChan)
    24  
    25  	outStream := streams.NewOut(NewWriterWithHook(cli.OutBuffer(), func(p []byte) {
    26  		writerHookChan <- struct{}{}
    27  	}))
    28  	cli.SetOut(outStream)
    29  
    30  	r, _, err := os.Pipe()
    31  	assert.NilError(t, err)
    32  	cli.SetIn(streams.NewIn(r))
    33  
    34  	go func() {
    35  		errChan <- cmd.ExecuteContext(ctx)
    36  	}()
    37  
    38  	writeCtx, writeCancel := context.WithTimeout(ctx, 100*time.Millisecond)
    39  	defer writeCancel()
    40  
    41  	// wait for the prompt to be ready
    42  	select {
    43  	case <-writeCtx.Done():
    44  		t.Fatalf("command %s did not write prompt to stdout", cmd.Name())
    45  	case <-writerHookChan:
    46  		// drain the channel for future buffer writes
    47  		go func() {
    48  			for {
    49  				select {
    50  				case <-ctx.Done():
    51  					return
    52  				case <-writerHookChan:
    53  				}
    54  			}
    55  		}()
    56  	}
    57  
    58  	assert.Check(t, cli.OutBuffer().Len() > 0)
    59  
    60  	// a small delay to ensure the plugin is prompting
    61  	time.Sleep(100 * time.Microsecond)
    62  
    63  	errCtx, errCancel := context.WithTimeout(ctx, 100*time.Millisecond)
    64  	defer errCancel()
    65  
    66  	// sigint and sigterm are caught by the prompt
    67  	// this allows us to gracefully exit the prompt with a 0 exit code
    68  	syscall.Kill(syscall.Getpid(), syscall.SIGINT)
    69  
    70  	select {
    71  	case <-errCtx.Done():
    72  		t.Logf("command stdout:\n%s\n", cli.OutBuffer().String())
    73  		t.Logf("command stderr:\n%s\n", cli.ErrBuffer().String())
    74  		t.Fatalf("command %s did not return after SIGINT", cmd.Name())
    75  	case err := <-errChan:
    76  		if assertFunc != nil {
    77  			assertFunc(t, err)
    78  		} else {
    79  			assert.NilError(t, err)
    80  			assert.Equal(t, cli.ErrBuffer().String(), "")
    81  		}
    82  	}
    83  }