From 5619626e169bd09a64abdb18e210107a1895dc39 Mon Sep 17 00:00:00 2001
From: Denys Kashkovskyi <dannie.k@me.com>
Date: Sat, 16 Apr 2022 12:13:55 +0000
Subject: [PATCH] Use a dedicated store for test results

---
 cmd/watch.go    | 12 ++----------
 core/printer.go | 20 ++++++++------------
 core/store.go   | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 22 deletions(-)
 create mode 100644 core/store.go

diff --git a/cmd/watch.go b/cmd/watch.go
index a354969..275b989 100644
--- a/cmd/watch.go
+++ b/cmd/watch.go
@@ -43,6 +43,7 @@ func runWatch(cmd *cobra.Command, args []string) {
 	resMap := sync.Map{}
 	printer := core.NewPrinter()
 	watcher := core.NewWatcher(&watchConfig)
+	store := core.NewStore()
 
 	go func() {
 		for {
@@ -51,16 +52,7 @@ func runWatch(cmd *cobra.Command, args []string) {
 		}
 	}()
 
-	watcher.Watch(func(res core.TestResult) {
-		if res.InProgress {
-			_, ok := resMap.Load(res.Id)
-			if !ok {
-				resMap.Store(res.Id, res)
-			}
-		} else {
-			resMap.Store(res.Id, res)
-		}
-	})
+	watcher.Watch(store.AddOrUpdate)
 }
 
 func init() {
diff --git a/core/printer.go b/core/printer.go
index 0b9c739..599c8b2 100644
--- a/core/printer.go
+++ b/core/printer.go
@@ -4,7 +4,6 @@ import (
 	"os"
 	"os/exec"
 	"runtime"
-	"sync"
 
 	"github.com/jedib0t/go-pretty/v6/table"
 )
@@ -42,19 +41,16 @@ func NewPrinter() Printer {
 	return Printer{clearFns: clear, t: t}
 }
 
-func (p *Printer) ToTable(results *sync.Map) {
+func (p *Printer) ToTable(results *Store) {
 	p.Clear()
 	p.t.ResetRows()
-	results.Range(func(k any, r interface{}) bool {
-		testResult, ok := r.(TestResult)
-		if ok {
-			p.t.AppendRow(table.Row{
-				k,
-				testResult.Tcp,
-				testResult.HttpStatus,
-				testResult.Duration,
-			})
-		}
+	results.ForEach(func(r TestResult) bool {
+		p.t.AppendRow(table.Row{
+			r.Id,
+			r.Tcp,
+			r.HttpStatus,
+			r.Duration,
+		})
 		return true
 	})
 	p.t.SortBy([]table.SortBy{{Name: "Address"}})
diff --git a/core/store.go b/core/store.go
new file mode 100644
index 0000000..81168b8
--- /dev/null
+++ b/core/store.go
@@ -0,0 +1,45 @@
+package core
+
+import "sync"
+
+type Store struct {
+	results sync.Map
+}
+
+func NewStore() Store {
+	syncMap := sync.Map{}
+
+	return Store{
+		results: syncMap,
+	}
+}
+
+func (s *Store) AddOrUpdate(res TestResult) {
+	if res.InProgress {
+		existing, ok := s.results.Load(res.Id)
+		if !ok {
+			s.results.Store(res.Id, res)
+		} else {
+			prev := existing.(TestResult)
+			s.results.Store(res.Id, TestResult{
+				Id:         prev.Id,
+				InProgress: true,
+				Tcp:        prev.Tcp,
+				HttpStatus: prev.HttpStatus,
+				Duration:   prev.Duration,
+			})
+		}
+	} else {
+		s.results.Store(res.Id, res)
+	}
+}
+
+func (s *Store) Clear() {
+	s.results = sync.Map{}
+}
+
+func (s *Store) ForEach(f func(TestResult) bool) {
+	s.results.Range(func(key, value any) bool {
+		return f(value.(TestResult))
+	})
+}