EnvManager: fix race condition in test

Refactor EnvManager so it takes TaskManager as an input.
This way clients can interject created tasks and install listeners.

Change-Id: Ief332060aea9e98587b43d04d68a4640717d18cd
diff --git a/core/installer/cmd/env_manager.go b/core/installer/cmd/env_manager.go
index e6b8499..aa33de8 100644
--- a/core/installer/cmd/env_manager.go
+++ b/core/installer/cmd/env_manager.go
@@ -9,6 +9,7 @@
 	"github.com/giolekva/pcloud/core/installer/dns"
 	"github.com/giolekva/pcloud/core/installer/http"
 	"github.com/giolekva/pcloud/core/installer/soft"
+	"github.com/giolekva/pcloud/core/installer/tasks"
 	"github.com/giolekva/pcloud/core/installer/welcome"
 )
 
@@ -86,6 +87,7 @@
 		installer.NewFixedLengthRandomNameGenerator(4),
 		httpClient,
 		dns.NewClient(),
+		tasks.NewTaskMap(),
 	)
 	log.Printf("Starting server\n")
 	s.Start()
diff --git a/core/installer/tasks/manager.go b/core/installer/tasks/manager.go
new file mode 100644
index 0000000..3dddccb
--- /dev/null
+++ b/core/installer/tasks/manager.go
@@ -0,0 +1,34 @@
+package tasks
+
+import (
+	"fmt"
+)
+
+type TaskManager interface {
+	Add(name string, task Task) error
+	Get(name string) (Task, error)
+}
+
+type TaskMap struct {
+	t map[string]Task
+}
+
+func NewTaskMap() *TaskMap {
+	return &TaskMap{make(map[string]Task)}
+}
+
+func (m *TaskMap) Add(name string, task Task) error {
+	if _, ok := m.t[name]; ok {
+		return fmt.Errorf("already exists")
+	}
+	m.t[name] = task
+	return nil
+}
+
+func (m *TaskMap) Get(name string) (Task, error) {
+	if t, ok := m.t[name]; ok {
+		return t, nil
+	} else {
+		return nil, fmt.Errorf("does not exist")
+	}
+}
diff --git a/core/installer/welcome/env.go b/core/installer/welcome/env.go
index 388cba3..15be0eb 100644
--- a/core/installer/welcome/env.go
+++ b/core/installer/welcome/env.go
@@ -88,7 +88,7 @@
 	nameGenerator installer.NameGenerator
 	httpClient    phttp.Client
 	dnsClient     dns.Client
-	Tasks         map[string]tasks.Task
+	Tasks         tasks.TaskManager
 	envInfo       map[string]template.HTML
 	dns           map[string]installer.EnvDNS
 	dnsPublished  map[string]struct{}
@@ -104,6 +104,7 @@
 	nameGenerator installer.NameGenerator,
 	httpClient phttp.Client,
 	dnsClient dns.Client,
+	tm tasks.TaskManager,
 ) *EnvServer {
 	return &EnvServer{
 		port,
@@ -115,7 +116,7 @@
 		nameGenerator,
 		httpClient,
 		dnsClient,
-		make(map[string]tasks.Task),
+		tm,
 		make(map[string]template.HTML),
 		make(map[string]installer.EnvDNS),
 		make(map[string]struct{}),
@@ -141,9 +142,9 @@
 		http.Error(w, "Task key not provided", http.StatusBadRequest)
 		return
 	}
-	t, ok := s.Tasks[key]
-	if !ok {
-		http.Error(w, "Task not found", http.StatusBadRequest)
+	t, err := s.Tasks.Get(key)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
 	dnsRecords := ""
@@ -410,7 +411,10 @@
 		mgr,
 		infoUpdater,
 	)
-	s.Tasks[key] = t
+	if err := s.Tasks.Add(key, t); err != nil {
+		panic(err)
+	}
+
 	s.dns[key] = dns
 	go t.Start()
 	http.Redirect(w, r, fmt.Sprintf("/env/%s", key), http.StatusSeeOther)
diff --git a/core/installer/welcome/env_test.go b/core/installer/welcome/env_test.go
index 0d61af8..e4cee83 100644
--- a/core/installer/welcome/env_test.go
+++ b/core/installer/welcome/env_test.go
@@ -12,6 +12,7 @@
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/go-git/go-billy/v5"
 	"github.com/go-git/go-billy/v5/memfs"
@@ -21,6 +22,7 @@
 
 	"github.com/giolekva/pcloud/core/installer"
 	"github.com/giolekva/pcloud/core/installer/soft"
+	"github.com/giolekva/pcloud/core/installer/tasks"
 )
 
 type fakeNSCreator struct {
@@ -184,6 +186,24 @@
 	return []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("2.2.2.2")}, nil
 }
 
+type onDoneTaskMap struct {
+	m      tasks.TaskManager
+	onDone tasks.TaskDoneListener
+}
+
+func (m *onDoneTaskMap) Add(name string, task tasks.Task) error {
+	if err := m.m.Add(name, task); err != nil {
+		return err
+	} else {
+		task.OnDone(m.onDone)
+		return nil
+	}
+}
+
+func (m *onDoneTaskMap) Get(name string) (tasks.Task, error) {
+	return m.m.Get(name)
+}
+
 func TestCreateNewEnv(t *testing.T) {
 	apps := installer.NewInMemoryAppRepository(installer.CreateAllApps())
 	infraFS := memfs.New()
@@ -214,6 +234,16 @@
 	cg := fakeClientGetter{t, envFS}
 	httpClient := fakeHttpClient{t, make(map[string]int)}
 	dnsClient := fakeDnsClient{t, make(map[string]int)}
+	var done sync.WaitGroup
+	done.Add(1)
+	var taskErr error
+	tm := &onDoneTaskMap{
+		tasks.NewTaskMap(),
+		func(err error) {
+			taskErr = err
+			done.Done()
+		},
+	}
 	s := NewEnvServer(
 		8181,
 		fakeSoftServeClient{t, envFS},
@@ -224,8 +254,10 @@
 		fixedNameGenerator{},
 		httpClient,
 		dnsClient,
+		tm,
 	)
 	go s.Start()
+	time.Sleep(1 * time.Second) // Let server start
 	req := createEnvReq{
 		Name:           "test",
 		ContactEmail:   "test@test.t",
@@ -238,13 +270,6 @@
 		t.Fatal(err)
 	}
 	resp, err := http.Post("http://localhost:8181/", "application/json", &buf)
-	var done sync.WaitGroup
-	done.Add(1)
-	var taskErr error
-	s.Tasks["test"].OnDone(func(err error) {
-		taskErr = err
-		done.Done()
-	})
 	if err != nil {
 		t.Fatal(err)
 	}