go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/buildproxywrap/relay_test.go (about)

     1  // Copyright 2023 The Fuchsia Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can
     3  // found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"os"
    11  	"strings"
    12  	"testing"
    13  	"time"
    14  )
    15  
    16  func noWait() error {
    17  	return nil
    18  }
    19  
    20  func TestRunInterruptOnCancel(t *testing.T) {
    21  	t.Run("kill cat", func(t *testing.T) {
    22  		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(time.Millisecond*50))
    23  		defer cancel()
    24  		// 'cat' hangs indefinitely waiting for stdin.  Kill the cat.
    25  		err := runInterruptOnCancel(ctx, "[cat] ", noWait, "cat")
    26  		if err != nil {
    27  			t.Errorf("test KillCat: expected no error.  Got %v", err)
    28  		}
    29  	})
    30  
    31  	t.Run("normal exit", func(t *testing.T) {
    32  		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(time.Second*5))
    33  		defer cancel()
    34  		// 'echo' will complete right away.
    35  		err := runInterruptOnCancel(ctx, "[echo] ", noWait, "echo", "hello")
    36  		if err != nil {
    37  			t.Errorf("test NormalExit: expected no error.  Got %v", err)
    38  		}
    39  	})
    40  
    41  	t.Run("nonzero exit", func(t *testing.T) {
    42  		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(time.Second*5))
    43  		defer cancel()
    44  		// 'false' will fail right away with code 1.
    45  		err := runInterruptOnCancel(ctx, "[false] ", noWait, "false")
    46  		if err == nil {
    47  			t.Errorf("test NonzeroExit: expected error.  Got nil")
    48  		}
    49  	})
    50  }
    51  
    52  const (
    53  	fakeSocat = "./socket_sink.py"
    54  )
    55  
    56  func TestMultiRelayWrap(t *testing.T) {
    57  	t.Run("wrapped command exit zero", func(t *testing.T) {
    58  		fakeRelay := socketRelay{
    59  			Name:             "foo",
    60  			SocketFileName:   "foo.sock",
    61  			SocketPathEnvVar: "FOO_socket_path",
    62  			ServerAddr:       "foo-service.googleapis.com:443",
    63  		}
    64  
    65  		ctx := context.Background()
    66  		tempDir, dirErr := os.MkdirTemp("", "test_wrap_cmd_zero")
    67  		if dirErr != nil {
    68  			t.Errorf("TestMultiRelayWrap (zero): expected mkdir success.  Got %v", dirErr)
    69  		}
    70  		defer os.RemoveAll(tempDir)
    71  
    72  		// We don't really need a real 'socat' binary to test.
    73  		err := multiRelayWrap(ctx, []*socketRelay{&fakeRelay}, tempDir, fakeSocat, func(env []string) error {
    74  			e := env[0]
    75  			if !strings.HasPrefix(e, "FOO_socket_path=") {
    76  				t.Errorf("TestMultiRelayWrap (zero): expected \"%q\" to start with \"%s\"", e, "FOO_socket_path=")
    77  			}
    78  			if !strings.HasSuffix(e, "foo.sock") {
    79  				t.Errorf("TestMultiRelayWrap (zero): expected \"%q\" to end with \"%s\"", e, "foo.sock")
    80  			}
    81  			return nil
    82  		})
    83  		if err != nil {
    84  			t.Errorf("TestMultiRelayWrap (zero): expected no error.  Got %v", err)
    85  		}
    86  	})
    87  
    88  	t.Run("wrapped command exit nonzero", func(t *testing.T) {
    89  		fakeRelay := socketRelay{
    90  			Name:             "bar",
    91  			SocketFileName:   "bar.sock",
    92  			SocketPathEnvVar: "BAR_socket_path",
    93  			ServerAddr:       "bar-service.googleapis.com:443",
    94  		}
    95  
    96  		ctx := context.Background()
    97  		tempDir, dirErr := os.MkdirTemp("", "test_wrap_cmd_nonzero")
    98  		if dirErr != nil {
    99  			t.Errorf("TestMultiRelayWrap (nonzero): expected mkdir success.  Got %v", dirErr)
   100  		}
   101  		defer os.RemoveAll(tempDir)
   102  
   103  		// 'cat' will run until interrupted
   104  		err := multiRelayWrap(ctx, []*socketRelay{&fakeRelay}, tempDir, fakeSocat, func(env []string) error {
   105  			e := env[0]
   106  			if !strings.HasPrefix(e, "BAR_socket_path=") {
   107  				t.Errorf("TestMultiRelayWrap (nonzero): expected \"%q\" to start with \"%s\"", e, "BAR_socket_path=")
   108  			}
   109  			if !strings.HasSuffix(e, "bar.sock") {
   110  				t.Errorf("TestMultiRelayWrap (nonzero): expected \"%q\" to end with \"%s\"", e, "foo.sock")
   111  			}
   112  			// mimic nonzero exit
   113  			return fmt.Errorf("exit 1")
   114  		})
   115  		if err == nil {
   116  			t.Errorf("TestMultiRelayWrap (nonzero): expected error.  Got nil")
   117  		}
   118  	})
   119  }