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

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"io/ioutil"
     7  	"testing"
     8  )
     9  
    10  func TestProxyHandlerUnaryRPC(t *testing.T) {
    11  	ctx := context.Background()
    12  
    13  	backmux := NewRespondMux()
    14  	backmux.Handle("simple", HandlerFunc(func(r Responder, c *Call) {
    15  		r.Return("simple")
    16  	}))
    17  
    18  	backend, _ := newTestPair(backmux)
    19  	defer backend.Close()
    20  
    21  	frontmux := NewRespondMux()
    22  	frontmux.Handle("", ProxyHandler(backend))
    23  
    24  	client, _ := newTestPair(frontmux)
    25  	defer client.Close()
    26  
    27  	var out interface{}
    28  	_, err := client.Call(ctx, "simple", nil, &out)
    29  	fatal(t, err)
    30  	if out != "simple" {
    31  		t.Fatal("unexpected return:", out)
    32  	}
    33  }
    34  
    35  func TestProxyHandlerBytestream(t *testing.T) {
    36  	ctx := context.Background()
    37  
    38  	backmux := NewRespondMux()
    39  	backmux.Handle("echo", HandlerFunc(func(r Responder, c *Call) {
    40  		c.Receive(nil)
    41  		ch, err := r.Continue(nil)
    42  		fatal(t, err)
    43  		io.Copy(ch, ch)
    44  		ch.Close()
    45  	}))
    46  
    47  	backend, _ := newTestPair(backmux)
    48  	defer backend.Close()
    49  
    50  	frontmux := NewRespondMux()
    51  	frontmux.Handle("", ProxyHandler(backend))
    52  
    53  	client, _ := newTestPair(frontmux)
    54  	defer client.Close()
    55  
    56  	resp, err := client.Call(ctx, "echo", nil, nil)
    57  	fatal(t, err)
    58  	_, err = io.WriteString(resp.Channel, "Hello world")
    59  	fatal(t, err)
    60  	fatal(t, resp.Channel.CloseWrite())
    61  	b, err := ioutil.ReadAll(resp.Channel)
    62  	fatal(t, err)
    63  	if string(b) != "Hello world" {
    64  		t.Fatal("unexpected return data:", string(b))
    65  	}
    66  }