testutil.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. package testutil
  2. import (
  3. "io"
  4. "net"
  5. "time"
  6. "github.com/Sirupsen/logrus"
  7. "golang.org/x/net/context"
  8. )
  9. // Handler is function called to handle incoming connection
  10. type Handler func(ctx context.Context, conn net.Conn, meta map[string][]string) error
  11. // Dialer is a function for dialing an outgoing connection
  12. type Dialer func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error)
  13. // TestStream creates an in memory session dialer for a handler function
  14. func TestStream(handler Handler) Dialer {
  15. s1, s2 := sockPair()
  16. return func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error) {
  17. go func() {
  18. err := handler(context.TODO(), s1, meta)
  19. if err != nil {
  20. logrus.Error(err)
  21. }
  22. s1.Close()
  23. }()
  24. return s2, nil
  25. }
  26. }
  27. func sockPair() (*sock, *sock) {
  28. pr1, pw1 := io.Pipe()
  29. pr2, pw2 := io.Pipe()
  30. return &sock{pw1, pr2, pw1}, &sock{pw2, pr1, pw2}
  31. }
  32. type sock struct {
  33. io.Writer
  34. io.Reader
  35. io.Closer
  36. }
  37. func (s *sock) LocalAddr() net.Addr {
  38. return dummyAddr{}
  39. }
  40. func (s *sock) RemoteAddr() net.Addr {
  41. return dummyAddr{}
  42. }
  43. func (s *sock) SetDeadline(t time.Time) error {
  44. return nil
  45. }
  46. func (s *sock) SetReadDeadline(t time.Time) error {
  47. return nil
  48. }
  49. func (s *sock) SetWriteDeadline(t time.Time) error {
  50. return nil
  51. }
  52. type dummyAddr struct {
  53. }
  54. func (d dummyAddr) Network() string {
  55. return "tcp"
  56. }
  57. func (d dummyAddr) String() string {
  58. return "localhost"
  59. }