github.com/Kong/go-pdk@v0.11.0/bridge/bridgetest/bridgetest.go (about)

     1  package bridgetest
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"io"
     7  	"net"
     8  	"testing"
     9  
    10  	"google.golang.org/protobuf/proto"
    11  )
    12  
    13  type MockStep struct {
    14  	Method string
    15  	Args   proto.Message
    16  	Ret    proto.Message
    17  }
    18  
    19  func readPbFrame(conn net.Conn) (data []byte, err error) {
    20  	var len uint32
    21  	err = binary.Read(conn, binary.LittleEndian, &len)
    22  	if err != nil {
    23  		return
    24  	}
    25  
    26  	data = make([]byte, len)
    27  
    28  	_, err = io.ReadFull(conn, data)
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  
    33  	return
    34  }
    35  
    36  func writePbFrame(conn net.Conn, data []byte) (err error) {
    37  	var len uint32 = uint32(len(data))
    38  	err = binary.Write(conn, binary.LittleEndian, len)
    39  	if err != nil {
    40  		return
    41  	}
    42  
    43  	if len > 0 {
    44  		_, err = conn.Write(data)
    45  	}
    46  
    47  	return
    48  }
    49  
    50  func Mock(t *testing.T, s []MockStep) net.Conn {
    51  	conA, conB := net.Pipe()
    52  
    53  	go func() {
    54  		for i, stp := range s {
    55  			d, err := readPbFrame(conB)
    56  			if err != nil {
    57  				t.Errorf("step %d readPbFrame(method): %s", i, err)
    58  				break
    59  			}
    60  			if !bytes.Equal([]byte(stp.Method), d) {
    61  				t.Errorf("step %d, expected method %v, found %v", i, []byte(stp.Method), d)
    62  				break
    63  			}
    64  
    65  			d, err = readPbFrame(conB)
    66  			if err != nil {
    67  				t.Errorf("step %d, readPbFrame(args): %s", i, err)
    68  				break
    69  			}
    70  
    71  			if stp.Args != nil {
    72  				args_d, err := proto.Marshal(stp.Args)
    73  				if err != nil {
    74  					t.Errorf("step %d, Marshal(args): %s", i, err)
    75  					break
    76  				}
    77  
    78  				if !bytes.Equal(args_d, d) {
    79  					t.Errorf("step %d, expected %v(%v), received %v", i, stp.Args, args_d, d)
    80  					break
    81  				}
    82  			}
    83  
    84  			if stp.Ret != nil {
    85  				ret_enc, err := proto.Marshal(stp.Ret)
    86  				if err != nil {
    87  					t.Errorf("step %d, Marshal(ret): %s", i, err)
    88  					break
    89  				}
    90  
    91  				err = writePbFrame(conB, ret_enc)
    92  				if err != nil {
    93  					t.Errorf("step %d, writePbFrame(ret): %s", i, err)
    94  					break
    95  				}
    96  			} else {
    97  				err = writePbFrame(conB, []byte{})
    98  				if err != nil {
    99  					t.Errorf("step %d, writePbFrame(ret): %s", i, err)
   100  					break
   101  				}
   102  			}
   103  		}
   104  		conB.Close()
   105  	}()
   106  	return conA
   107  }
   108  
   109  type mockEnvironment interface {
   110  	Handle(method string, args_d []byte) []byte
   111  	Errorf(format string, args ...interface{})
   112  	IsRunning() bool
   113  	SubscribeStatusChange(ch chan<- string)
   114  }
   115  
   116  func MockFunc(e mockEnvironment) net.Conn {
   117  	conA, conB := net.Pipe()
   118  
   119  	statusCh := make(chan string, 1)
   120  	e.SubscribeStatusChange(statusCh)
   121  
   122  	go func() {
   123  		for {
   124  			d, err := readPbFrame(conB)
   125  			if err != nil {
   126  				e.Errorf("Can't read method name")
   127  				break
   128  			}
   129  			method := string(d)
   130  
   131  			d, err = readPbFrame(conB)
   132  			if err != nil {
   133  				e.Errorf("Can't read method \"%v\" arguments", method)
   134  				break
   135  			}
   136  
   137  			d = e.Handle(method, d)
   138  
   139  			err = writePbFrame(conB, d)
   140  			if err != nil {
   141  				e.Errorf("Can't write back return values")
   142  				break
   143  			}
   144  
   145  			select {
   146  			case msg := <-statusCh:
   147  				if msg == "finished" {
   148  					return
   149  				}
   150  			default: // do nothing
   151  			}
   152  		}
   153  	}()
   154  	return conA
   155  }