tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/mux/transport_test.go (about)

     1  package mux
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"io/ioutil"
     8  	"path"
     9  	"strings"
    10  	"testing"
    11  )
    12  
    13  func testExchange(t *testing.T, sess Session) {
    14  	var err error
    15  	var ch Channel
    16  	t.Run("session accept", func(t *testing.T) {
    17  		ch, err = sess.Accept()
    18  		fatal(err, t)
    19  	})
    20  
    21  	t.Run("channel write", func(t *testing.T) {
    22  		_, err = ch.Write([]byte("Hello world"))
    23  		fatal(err, t)
    24  		err = ch.Close()
    25  		fatal(err, t)
    26  	})
    27  
    28  	t.Run("session open", func(t *testing.T) {
    29  		ch, err = sess.Open(context.Background())
    30  		fatal(err, t)
    31  	})
    32  
    33  	var b []byte
    34  	t.Run("channel read", func(t *testing.T) {
    35  		b, err = ioutil.ReadAll(ch)
    36  		fatal(err, t)
    37  		err = ch.Close()
    38  		fatal(err, t)
    39  	})
    40  
    41  	if !bytes.Equal(b, []byte("Hello world")) {
    42  		t.Fatalf("unexpected bytes: %s", b)
    43  	}
    44  
    45  	t.Run("session close", func(t *testing.T) {
    46  		err = sess.Close()
    47  		fatal(err, t)
    48  	})
    49  }
    50  
    51  func startListener(t *testing.T, l Listener) {
    52  	t.Helper()
    53  
    54  	t.Cleanup(func() {
    55  		fatal(l.Close(), t)
    56  	})
    57  
    58  	go func() {
    59  		sess, err := l.Accept()
    60  		fatal(err, t)
    61  		t.Cleanup(func() {
    62  			// Synchronizes cleanup, waiting for the client to disconnect before
    63  			// closing the stream. This prevents errors in the Pipe-based test with
    64  			// closing one end of the pipe before the other has read the data.
    65  			// Registering as a test cleanup function also avoids a race condition
    66  			// with the test exiting before closing the session.
    67  			if err := sess.Wait(); err != io.EOF {
    68  				// windows test environments make wait return funky errors
    69  				if !strings.Contains(err.Error(), "wsarecv") {
    70  					t.Errorf("Wait returned unexpected error: %v", err)
    71  				}
    72  			}
    73  			err = sess.Close()
    74  			fatal(err, t)
    75  		})
    76  
    77  		ch, err := sess.Open(context.Background())
    78  		fatal(err, t)
    79  		b, err := ioutil.ReadAll(ch)
    80  		fatal(err, t)
    81  		ch.Close()
    82  
    83  		ch, err = sess.Accept()
    84  		_, err = ch.Write(b)
    85  		fatal(err, t)
    86  		err = ch.CloseWrite()
    87  		fatal(err, t)
    88  	}()
    89  }
    90  
    91  func TestTCP(t *testing.T) {
    92  	l, err := ListenTCP("127.0.0.1:0")
    93  	fatal(err, t)
    94  	startListener(t, l)
    95  
    96  	sess, err := DialTCP(l.Addr().String())
    97  	fatal(err, t)
    98  	testExchange(t, sess)
    99  }
   100  
   101  func TestUnix(t *testing.T) {
   102  	tmp := t.TempDir()
   103  	sockPath := path.Join(tmp, "qmux.sock")
   104  	l, err := ListenUnix(sockPath)
   105  	fatal(err, t)
   106  	startListener(t, l)
   107  
   108  	sess, err := DialUnix(sockPath)
   109  	fatal(err, t)
   110  	testExchange(t, sess)
   111  }
   112  
   113  func TestIO(t *testing.T) {
   114  	pr1, pw1 := io.Pipe()
   115  	pr2, pw2 := io.Pipe()
   116  
   117  	l, err := ListenIO(pw1, pr2)
   118  	fatal(err, t)
   119  	startListener(t, l)
   120  
   121  	sess, err := DialIO(pw2, pr1)
   122  	fatal(err, t)
   123  	testExchange(t, sess)
   124  }
   125  
   126  // func TestWS(t *testing.T) {
   127  // 	l, err := ListenWS("127.0.0.1:0")
   128  // 	fatal(err, t)
   129  // 	startListener(t, l)
   130  
   131  // 	sess, err := DialWS(l.Addr().String())
   132  // 	fatal(err, t)
   133  // 	testExchange(t, sess)
   134  // }