3 files changed, 75 insertions, 23 deletions
diff --git a/p2p/simulations/adapters/exec.go b/p2p/simulations/adapters/exec.go
index a566fb27d..f381c1159 100644
--- a/p2p/simulations/adapters/exec.go
+++ b/p2p/simulations/adapters/exec.go
@@ -17,7 +17,6 @@
 package adapters
 
 import (
-	"bufio"
 	"context"
 	"crypto/ecdsa"
 	"encoding/json"
@@ -29,7 +28,6 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
-	"regexp"
 	"strings"
 	"sync"
 	"syscall"
@@ -150,10 +148,6 @@ func (n *ExecNode) Client() (*rpc.Client, error) {
 	return n.client, nil
 }
 
-// wsAddrPattern is a regex used to read the WebSocket address from the node's
-// log
-var wsAddrPattern = regexp.MustCompile(`ws://[\d.:]+`)
-
 // Start exec's the node passing the ID and service as command line arguments
 // and the node config encoded as JSON in the _P2P_NODE_CONFIG environment
 // variable
@@ -196,23 +190,9 @@ func (n *ExecNode) Start(snapshots map[string][]byte) (err error) {
 	n.Cmd = cmd
 
 	// read the WebSocket address from the stderr logs
-	var wsAddr string
-	wsAddrC := make(chan string)
-	go func() {
-		s := bufio.NewScanner(stderrR)
-		for s.Scan() {
-			if strings.Contains(s.Text(), "WebSocket endpoint opened:") {
-				wsAddrC <- wsAddrPattern.FindString(s.Text())
-			}
-		}
-	}()
-	select {
-	case wsAddr = <-wsAddrC:
-		if wsAddr == "" {
-			return errors.New("failed to read WebSocket address from stderr")
-		}
-	case <-time.After(10 * time.Second):
-		return errors.New("timed out waiting for WebSocket address on stderr")
+	wsAddr, err := findWSAddr(stderrR, 10*time.Second)
+	if err != nil {
+		return fmt.Errorf("error getting WebSocket address: %s", err)
 	}
 
 	// create the RPC client and load the node info
diff --git a/p2p/simulations/adapters/ws.go b/p2p/simulations/adapters/ws.go
new file mode 100644
index 000000000..979a21709
--- /dev/null
+++ b/p2p/simulations/adapters/ws.go
@@ -0,0 +1,51 @@
+package adapters
+
+import (
+	"bufio"
+	"errors"
+	"io"
+	"regexp"
+	"strings"
+	"time"
+)
+
+// wsAddrPattern is a regex used to read the WebSocket address from the node's
+// log
+var wsAddrPattern = regexp.MustCompile(`ws://[\d.:]+`)
+
+func matchWSAddr(str string) (string, bool) {
+	if !strings.Contains(str, "WebSocket endpoint opened") {
+		return "", false
+	}
+
+	return wsAddrPattern.FindString(str), true
+}
+
+// findWSAddr scans through reader r, looking for the log entry with
+// WebSocket address information.
+func findWSAddr(r io.Reader, timeout time.Duration) (string, error) {
+	ch := make(chan string)
+
+	go func() {
+		s := bufio.NewScanner(r)
+		for s.Scan() {
+			addr, ok := matchWSAddr(s.Text())
+			if ok {
+				ch <- addr
+			}
+		}
+		close(ch)
+	}()
+
+	var wsAddr string
+	select {
+	case wsAddr = <-ch:
+		if wsAddr == "" {
+			return "", errors.New("empty result")
+		}
+	case <-time.After(timeout):
+		return "", errors.New("timed out")
+	}
+
+	return wsAddr, nil
+}
diff --git a/p2p/simulations/adapters/ws_test.go b/p2p/simulations/adapters/ws_test.go
new file mode 100644
index 000000000..0bb9ed2b2
--- /dev/null
+++ b/p2p/simulations/adapters/ws_test.go
@@ -0,0 +1,21 @@
+package adapters
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+func TestFindWSAddr(t *testing.T) {
+	line := `t=2018-05-02T19:00:45+0200 lvl=info msg="WebSocket endpoint opened"  node.id=26c65a606d1125a44695bc08573190d047152b6b9a776ccbbe593e90f91444d9c1ebdadac6a775ad9fdd0923468a1d698ed3a842c1fb89c1bc0f9d4801f8c39c url=ws://127.0.0.1:59975`
+	buf := bytes.NewBufferString(line)
+	got, err := findWSAddr(buf, 10*time.Second)
+	if err != nil {
+		t.Fatalf("Failed to find addr: %v", err)
+	}
+	expected := `ws://127.0.0.1:59975`
+
+	if got != expected {
+		t.Fatalf("Expected to get '%s', but got '%s'", expected, got)
+	}
+}