EnvManager: Fix infrastructure service wait addresses

Change-Id: I525fd989ee249b725023c7f118a745fb60b4c957
diff --git a/core/installer/tasks/infra.go b/core/installer/tasks/infra.go
index 130a89a..f855adf 100644
--- a/core/installer/tasks/infra.go
+++ b/core/installer/tasks/infra.go
@@ -265,11 +265,17 @@
 		}
 		return nil
 	})
+	var addr string
+	if env.PrivateDomain != "" {
+		addr = fmt.Sprintf("https://memberships.%s", env.PrivateDomain)
+	} else {
+		addr = fmt.Sprintf("https://memberships.%s", env.Domain)
+	}
 	return newSequentialParentTask(
 		"Group membership",
 		false,
 		&t,
-		waitForAddr(st.httpClient, fmt.Sprintf("https://memberships.p.%s", env.Domain)),
+		waitForAddr(st.httpClient, addr),
 	)
 }
 
@@ -374,7 +380,7 @@
 }
 
 func SetupAppStore(env installer.EnvConfig, st *state) Task {
-	t := newLeafTask("Application marketplace", func() error {
+	t := newLeafTask("Setup", func() error {
 		user := fmt.Sprintf("%s-appmanager", env.Id)
 		keys, err := installer.NewSSHKeyPair(user)
 		if err != nil {
@@ -407,7 +413,18 @@
 		}
 		return nil
 	})
-	return &t
+	var addr string
+	if env.PrivateDomain != "" {
+		addr = fmt.Sprintf("https://apps.%s", env.PrivateDomain)
+	} else {
+		addr = fmt.Sprintf("https://apps.%s", env.Domain)
+	}
+	return newSequentialParentTask(
+		"Application marketplace",
+		false,
+		&t,
+		waitForAddr(st.httpClient, addr),
+	)
 }
 
 // TODO(gio-dns): remove
diff --git a/core/installer/welcome/env_test.go b/core/installer/welcome/env_test.go
index 414fac8..1ab6b39 100644
--- a/core/installer/welcome/env_test.go
+++ b/core/installer/welcome/env_test.go
@@ -329,6 +329,7 @@
 		t.Fatal(taskErr)
 	}
 	expected := []string{
+		"https://apps.p.test.t",
 		"https://accounts-ui.test.t",
 		"https://welcome.test.t",
 		"https://memberships.p.test.t",
@@ -340,7 +341,7 @@
 			t.Fatal(httpClient.counts)
 		}
 	}
-	if len(httpClient.counts) != 5 {
+	if len(httpClient.counts) != 6 {
 		t.Fatal(httpClient.counts)
 	}
 }