github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/rpc/dispatch_test.go (about)

     1  // Copyright 2012, 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package rpc_test
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"net/http"
    10  	"net/http/httptest"
    11  
    12  	"github.com/gorilla/websocket"
    13  	"github.com/juju/loggo"
    14  	jc "github.com/juju/testing/checkers"
    15  	gc "gopkg.in/check.v1"
    16  
    17  	"github.com/juju/juju/rpc"
    18  	"github.com/juju/juju/rpc/jsoncodec"
    19  	"github.com/juju/juju/testing"
    20  )
    21  
    22  type dispatchSuite struct {
    23  	testing.BaseSuite
    24  
    25  	server     *httptest.Server
    26  	serverAddr string
    27  	ready      chan struct{}
    28  }
    29  
    30  var _ = gc.Suite(&dispatchSuite{})
    31  
    32  func (s *dispatchSuite) SetUpSuite(c *gc.C) {
    33  	s.BaseSuite.SetUpSuite(c)
    34  	rpcServer := func(ws *websocket.Conn) {
    35  		codec := jsoncodec.NewWebsocket(ws)
    36  		conn := rpc.NewConn(codec, nil)
    37  
    38  		conn.Serve(&DispatchRoot{}, nil, nil)
    39  		conn.Start(context.Background())
    40  
    41  		<-conn.Dead()
    42  	}
    43  	http.Handle("/rpc", websocketHandler(rpcServer))
    44  	s.server = httptest.NewServer(nil)
    45  	s.serverAddr = s.server.Listener.Addr().String()
    46  	s.ready = make(chan struct{}, 1)
    47  	s.AddCleanup(func(*gc.C) {
    48  		s.server.Close()
    49  	})
    50  }
    51  
    52  var wsUpgrader = &websocket.Upgrader{
    53  	CheckOrigin: func(*http.Request) bool {
    54  		return true
    55  	},
    56  }
    57  
    58  func websocketHandler(f func(*websocket.Conn)) http.Handler {
    59  	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    60  		c, err := wsUpgrader.Upgrade(w, req, nil)
    61  		if err == nil {
    62  			f(c)
    63  		}
    64  	})
    65  }
    66  
    67  func (s *dispatchSuite) SetUpTest(c *gc.C) {
    68  	s.BaseSuite.SetUpTest(c)
    69  	loggo.GetLogger("juju.rpc").SetLogLevel(loggo.TRACE)
    70  }
    71  
    72  func (s *dispatchSuite) TestWSWithoutParamsV0(c *gc.C) {
    73  	resp := s.request(c, `{"RequestId":1,"Type": "DispatchDummy","Id": "without","Request":"DoSomething"}`)
    74  	s.assertResponse(c, resp, `{"RequestId":1,"Response":{}}`)
    75  }
    76  
    77  func (s *dispatchSuite) TestWSWithParamsV0(c *gc.C) {
    78  	resp := s.request(c, `{"RequestId":2,"Type": "DispatchDummy","Id": "with","Request":"DoSomething", "Params": {}}`)
    79  	s.assertResponse(c, resp, `{"RequestId":2,"Response":{}}`)
    80  }
    81  
    82  func (s *dispatchSuite) TestWSWithoutParamsV1(c *gc.C) {
    83  	resp := s.request(c, `{"request-id":1,"type": "DispatchDummy","id": "without","request":"DoSomething"}`)
    84  	s.assertResponse(c, resp, `{"request-id":1,"response":{}}`)
    85  }
    86  
    87  func (s *dispatchSuite) TestWSWithParamsV1(c *gc.C) {
    88  	resp := s.request(c, `{"request-id":2,"type": "DispatchDummy","id": "with","request":"DoSomething", "params": {}}`)
    89  	s.assertResponse(c, resp, `{"request-id":2,"response":{}}`)
    90  }
    91  
    92  func (s *dispatchSuite) assertResponse(c *gc.C, obtained, expected string) {
    93  	c.Assert(obtained, gc.Equals, expected+"\n")
    94  }
    95  
    96  // request performs one request to the test server via websockets.
    97  func (s *dispatchSuite) request(c *gc.C, req string) string {
    98  	url := fmt.Sprintf("ws://%s/rpc", s.serverAddr)
    99  	ws, _, err := websocket.DefaultDialer.Dial(url, http.Header{
   100  		"Origin": {"http://localhost"},
   101  	})
   102  	c.Assert(err, jc.ErrorIsNil)
   103  
   104  	reqdata := []byte(req)
   105  	err = ws.WriteMessage(websocket.TextMessage, reqdata)
   106  	c.Assert(err, jc.ErrorIsNil)
   107  
   108  	_, resp, err := ws.ReadMessage()
   109  	c.Assert(err, jc.ErrorIsNil)
   110  
   111  	err = ws.Close()
   112  	c.Assert(err, jc.ErrorIsNil)
   113  
   114  	return string(resp)
   115  }
   116  
   117  // DispatchRoot simulates the root for the test.
   118  type DispatchRoot struct{}
   119  
   120  func (*DispatchRoot) DispatchDummy(id string) (*DispatchDummy, error) {
   121  	return &DispatchDummy{}, nil
   122  }
   123  
   124  // DispatchDummy is the type to whish the request is dispatched.
   125  type DispatchDummy struct{}
   126  
   127  func (d *DispatchDummy) DoSomething() {}