From 709d49218a572a24f1be589ddcefc4ef7cd66162 Mon Sep 17 00:00:00 2001
From: Rafael Matias <rafael@skyle.net>
Date: Tue, 26 Mar 2019 18:48:09 +0100
Subject: [PATCH] cmd/swarm/swarm-smoke-pss: add msg timeout via flag

---
 cmd/swarm/swarm-smoke-pss/main.go | 25 ++++++++++++++++---------
 cmd/swarm/swarm-smoke-pss/pss.go  |  4 ++--
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/cmd/swarm/swarm-smoke-pss/main.go b/cmd/swarm/swarm-smoke-pss/main.go
index be958e2cea..dc4f28b7bd 100644
--- a/cmd/swarm/swarm-smoke-pss/main.go
+++ b/cmd/swarm/swarm-smoke-pss/main.go
@@ -37,14 +37,15 @@ var (
 )
 
 var (
-	allhosts        string
-	hosts           []string
-	inputSeed       int
-	wsPort          int
-	verbosity       int
-	timeout         int
-	pssMessageCount int
-	pssMessageSize  int
+	allhosts          string
+	hosts             []string
+	inputSeed         int
+	wsPort            int
+	verbosity         int
+	timeout           int
+	pssMessageTimeout int
+	pssMessageCount   int
+	pssMessageSize    int
 )
 
 func main() {
@@ -84,6 +85,12 @@ func main() {
 			Usage:       "timeout in seconds after which kill the process",
 			Destination: &timeout,
 		},
+		cli.IntFlag{
+			Name:        "msgtimeout",
+			Value:       1,
+			Usage:       "timeout in seconds until a message is received",
+			Destination: &pssMessageTimeout,
+		},
 		cli.IntFlag{
 			Name:        "msgcount",
 			Value:       10,
@@ -128,7 +135,7 @@ func main() {
 		{
 			Name:   "all",
 			Usage:  "send and receive raw, sym and asym messages across randome nodes",
-			Action: wrapCliCommand("raw", pssAllCheck),
+			Action: wrapCliCommand("all", pssAllCheck),
 		},
 	}
 
diff --git a/cmd/swarm/swarm-smoke-pss/pss.go b/cmd/swarm/swarm-smoke-pss/pss.go
index e4f47c3be6..6636a51234 100644
--- a/cmd/swarm/swarm-smoke-pss/pss.go
+++ b/cmd/swarm/swarm-smoke-pss/pss.go
@@ -303,8 +303,8 @@ func (s *pssSession) waitForMsg() error {
 			return fmt.Errorf("message mismatch. expected: %s got: %s", resJob.msg, res.Msg)
 		}
 		log.Debug("got msg", "job", hexutil.Encode([]byte(resIdx)))
-	case <-time.After(1 * time.Second):
-		return errors.New("msg timeout after 1 sec")
+	case <-time.After(time.Duration(pssMessageTimeout) * time.Second):
+		return fmt.Errorf("message timeout after %d sec", pssMessageTimeout)
 	}
 	return nil
 }