From a60f376846171dc052c1076fa1f1c6fd30bb928b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arturo=20Filast=C3=B2?= <arturo@filasto.net>
Date: Mon, 24 Sep 2018 15:57:28 +0200
Subject: [PATCH] Adjust the progress bars depending on the test type

---
 internal/cli/run/run.go |  3 ++-
 nettests/nettests.go    | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/internal/cli/run/run.go b/internal/cli/run/run.go
index d7d1c24..76afd2a 100644
--- a/internal/cli/run/run.go
+++ b/internal/cli/run/run.go
@@ -80,9 +80,10 @@ func init() {
 			return err
 		}
 
-		for _, nt := range group.Nettests {
+		for i, nt := range group.Nettests {
 			log.Debugf("Running test %T", nt)
 			ctl := nettests.NewController(nt, ctx, result)
+			ctl.SetNettestIndex(i, len(group.Nettests))
 			if err = nt.Run(ctl); err != nil {
 				log.WithError(err).Errorf("Failed to run %s", group.Label)
 				return err
diff --git a/nettests/nettests.go b/nettests/nettests.go
index 4ac5ef7..5092659 100644
--- a/nettests/nettests.go
+++ b/nettests/nettests.go
@@ -45,6 +45,8 @@ type Controller struct {
 	Ctx         *ooni.Context
 	res         *database.Result
 	nt          Nettest
+	ntCount     int
+	ntIndex     int
 	msmts       map[int64]*database.Measurement
 	msmtPath    string          // XXX maybe we can drop this and just use a temporary file
 	inputIdxMap map[int64]int64 // Used to map mk idx to database id
@@ -58,11 +60,21 @@ func getCaBundlePath() string {
 	return "/etc/ssl/cert.pem"
 }
 
+// SetInputIdxMap is used to set the mapping of index into input. This mapping
+// is used to reference, for example, a particular URL based on the index inside
+// of the input list and the index of it in the database.
 func (c *Controller) SetInputIdxMap(inputIdxMap map[int64]int64) error {
 	c.inputIdxMap = inputIdxMap
 	return nil
 }
 
+// SetNettestIndex is used to set the current nettest index and total nettest
+// count to compute a different progress percentage.
+func (c *Controller) SetNettestIndex(i, n int) {
+	c.ntCount = n
+	c.ntIndex = i
+}
+
 // Init should be called once to initialise the nettest
 func (c *Controller) Init(nt *mk.Nettest) error {
 	log.Debugf("Init: %v", nt)
@@ -157,7 +169,11 @@ func (c *Controller) Init(nt *mk.Nettest) error {
 
 	nt.On("status.progress", func(e mk.Event) {
 		log.Debugf(color.RedString(e.Key))
-		c.OnProgress(e.Value.Percentage, e.Value.Message)
+		perc := e.Value.Percentage
+		if c.ntCount > 0 {
+			perc = float64(c.ntIndex)/float64(c.ntCount) + perc/float64(c.ntCount)
+		}
+		c.OnProgress(perc, e.Value.Message)
 	})
 
 	nt.On("status.update.*", func(e mk.Event) {