Merge pull request #19 from getamis/feature/add-websocket-support

container, tests: add websocket support

container/ethereum.go

@@ -90,6 +90,7 @@ type ethereum struct {
 	dataDir     string
 	port        string
 	rpcPort     string
+	wsPort      string
 	hostName    string
 	containerID string
 
@@ -153,34 +154,51 @@ func (eth *ethereum) Init(genesisFile string) error {
 }
 
 func (eth *ethereum) Start() error {
+	exposedPorts := make(map[nat.Port]struct{})
+	portBindings := nat.PortMap{}
+
+	if eth.port != "" {
+		exposedPorts[nat.Port(eth.port)] = struct{}{}
+		portBindings[nat.Port(eth.port)] = []nat.PortBinding{
+			{
+				HostIP:   "0.0.0.0",
+				HostPort: eth.port,
+			},
+		}
+	}
+
+	if eth.rpcPort != "" {
+		exposedPorts[nat.Port(eth.rpcPort)] = struct{}{}
+		portBindings[nat.Port(eth.rpcPort)] = []nat.PortBinding{
+			{
+				HostIP:   "0.0.0.0",
+				HostPort: eth.rpcPort,
+			},
+		}
+	}
+
+	if eth.wsPort != "" {
+		exposedPorts[nat.Port(eth.wsPort)] = struct{}{}
+		portBindings[nat.Port(eth.wsPort)] = []nat.PortBinding{
+			{
+				HostIP:   "0.0.0.0",
+				HostPort: eth.wsPort,
+			},
+		}
+	}
+
 	resp, err := eth.client.ContainerCreate(context.Background(),
 		&container.Config{
-			Hostname: "geth-" + eth.hostName,
-			Image:    eth.Image(),
-			Cmd:      eth.flags,
-			ExposedPorts: map[nat.Port]struct{}{
-				nat.Port(eth.port):    {},
-				nat.Port(eth.rpcPort): {},
-			},
+			Hostname:     "geth-" + eth.hostName,
+			Image:        eth.Image(),
+			Cmd:          eth.flags,
+			ExposedPorts: exposedPorts,
 		},
 		&container.HostConfig{
 			Binds: []string{
 				eth.hostDataDir + ":" + eth.dataDir,
 			},
-			PortBindings: nat.PortMap{
-				nat.Port(eth.port): []nat.PortBinding{
-					{
-						HostIP:   "0.0.0.0",
-						HostPort: eth.port,
-					},
-				},
-				nat.Port(eth.rpcPort): []nat.PortBinding{
-					{
-						HostIP:   "0.0.0.0",
-						HostPort: eth.rpcPort,
-					},
-				},
-			},
+			PortBindings: portBindings,
 		}, nil, "")
 	if err != nil {
 		log.Printf("Failed to create container, err: %v", err)
@@ -259,7 +277,17 @@ func (eth *ethereum) Running() bool {
 }
 
 func (eth *ethereum) NewClient() *ethclient.Client {
-	client, err := ethclient.Dial("http://" + eth.Host() + ":" + eth.rpcPort)
+	var scheme, port string
+
+	if eth.rpcPort != "" {
+		scheme = "http://"
+		port = eth.rpcPort
+	}
+	if eth.wsPort != "" {
+		scheme = "ws://"
+		port = eth.wsPort
+	}
+	client, err := ethclient.Dial(scheme + eth.Host() + ":" + port)
 	if err != nil {
 		log.Printf("Failed to dial to geth, err: %v", err)
 		return nil

container/ethereum_test.go

@@ -40,10 +40,12 @@ func TestEthereumContainer(t *testing.T) {
 			HostDataDir(env.DataDir),
 			DataDir("/data"),
 			Port(fmt.Sprintf("%d", env.P2PPort)),
-			RPC(),
-			RPCAddress("0.0.0.0"),
-			RPCAPI("eth,net,web3,personal"),
-			RPCPort(fmt.Sprintf("%d", env.RpcPort)),
+			WebSocket(),
+			WebSocketAddress("0.0.0.0"),
+			WebSocketAPI("eth,net,web3,personal"),
+			WebSocketPort(fmt.Sprintf("%d", env.RpcPort)),
+			WebSocketOrigin("*"),
+			NoDiscover(),
 		)
 
 		err := geth.Init(filepath.Join(env.DataDir, core.GenesisFile))

container/options.go

@@ -169,6 +169,41 @@ func RPCPort(port string) Option {
 	}
 }
 
+func WebSocket() Option {
+	return func(eth *ethereum) {
+		eth.flags = append(eth.flags, "--"+utils.WSEnabledFlag.Name)
+	}
+}
+
+func WebSocketAddress(address string) Option {
+	return func(eth *ethereum) {
+		eth.flags = append(eth.flags, "--"+utils.WSListenAddrFlag.Name)
+		eth.flags = append(eth.flags, address)
+	}
+}
+
+func WebSocketAPI(apis string) Option {
+	return func(eth *ethereum) {
+		eth.flags = append(eth.flags, "--"+utils.WSApiFlag.Name)
+		eth.flags = append(eth.flags, apis)
+	}
+}
+
+func WebSocketPort(port string) Option {
+	return func(eth *ethereum) {
+		eth.flags = append(eth.flags, "--"+utils.WSPortFlag.Name)
+		eth.flags = append(eth.flags, port)
+		eth.wsPort = port
+	}
+}
+
+func WebSocketOrigin(origins string) Option {
+	return func(eth *ethereum) {
+		eth.flags = append(eth.flags, "--"+utils.WSAllowedOriginsFlag.Name)
+		eth.flags = append(eth.flags, origins)
+	}
+}
+
 func Verbosity(verbosity int) Option {
 	return func(eth *ethereum) {
 		eth.flags = append(eth.flags, "--verbosity")

tests/integration_test.go

@@ -54,12 +54,15 @@ var _ = Describe("4 validators Istanbul", func() {
 				container.HostDataDir(env.DataDir),
 				container.DataDir("/data"),
 				container.Port(fmt.Sprintf("%d", env.P2PPort)),
-				container.RPC(),
-				container.RPCAddress("0.0.0.0"),
-				container.RPCAPI("eth,net,web3,personal"),
-				container.RPCPort(fmt.Sprintf("%d", env.RpcPort)),
+				container.WebSocket(),
+				container.WebSocketAddress("0.0.0.0"),
+				container.WebSocketAPI("eth,net,web3,personal,miner"),
+				container.WebSocketPort(fmt.Sprintf("%d", env.RpcPort)),
+				container.WebSocketOrigin("*"),
 				container.NAT("any"),
 				container.NoDiscover(),
+				container.Etherbase("1a9afb711302c5f83b5902843d1c007a1a137632"),
+				container.Mine(),
 				container.Logging(true),
 			)