functional: make replace tests wait for N active units

After "fleetctl start --replace", it should wait until expected number
of active units are available. So use waitForNActiveUnits() after
running list-units. However, for "fleetctl {load,submit} --replace",
it should expect 0 active units. To distinguish that, create a new
wrapper waitForActiveUnitsReplaceCmds().

functional/unit_action_test.go

@@ -313,6 +313,9 @@ func replaceUnitCommon(cmd string) error {
 	if len(units) != 1 {
 		return fmt.Errorf("Did not find 1 unit in cluster: \n%s", stdout)
 	}
+	if err := waitForActiveUnitsReplaceCmds(cluster, m, cmd, 1); err != nil {
+		return err
+	}
 
 	// replace the unit and assert it shows up
 	err = genNewFleetService(tmpHelloService, fxtHelloService, "sleep 2", "sleep 1")
@@ -330,6 +333,9 @@ func replaceUnitCommon(cmd string) error {
 	if len(units) != 1 {
 		return fmt.Errorf("Did not find 1 unit in cluster: \n%s", stdout)
 	}
+	if err := waitForActiveUnitsReplaceCmds(cluster, m, cmd, 1); err != nil {
+		return err
+	}
 	os.Remove(tmpHelloService)
 
 	if err := destroyUnitRetrying(cluster, m, fxtHelloService); err != nil {
@@ -397,6 +403,9 @@ func replaceUnitMultiple(cmd string, n int) error {
 		if len(units) != i {
 			return fmt.Errorf("Did not find %d units in cluster: \n%s", i, stdout)
 		}
+		if err := waitForActiveUnitsReplaceCmds(cluster, m, cmd, i); err != nil {
+			return err
+		}
 
 		// generate a new service derived by fixtures, and store it under /tmp
 		err = genNewFleetService(curHelloService, fxtHelloService, "sleep 2", "sleep 1")
@@ -420,6 +429,9 @@ func replaceUnitMultiple(cmd string, n int) error {
 		if len(units) != n {
 			return fmt.Errorf("Did not find %d units in cluster: \n%s", n, stdout)
 		}
+		if err := waitForActiveUnitsReplaceCmds(cluster, m, cmd, i); err != nil {
+			return err
+		}
 	}
 
 	// clean up temp services under /tmp
@@ -472,6 +484,35 @@ func genNewFleetService(newFile, oldFile, newVal, oldVal string) error {
 	return nil
 }
 
+// waitForActiveUnitsReplaceCmds() is a wrapper for waiting for N active units.
+// The expected number of active units are given as a parameter "count".
+// If cmd is "start", it expects that "count" active units are active.
+// Otherwise, for "load" or "submit", it expects no active unit.
+func waitForActiveUnitsReplaceCmds(cluster platform.Cluster, m platform.Member, cmd string, count int) error {
+	if cmd == "start" {
+		units, err := cluster.WaitForNActiveUnits(m, count)
+		if err != nil {
+			fmt.Errorf("%v", err)
+		}
+		_, found := units["hello.service"]
+		if len(units) != count || !found {
+			fmt.Errorf("Expected hello.service to be sole active unit, got %v", units)
+		}
+	} else {
+		// cmd is "load" or "submit", then there's no active unit
+		units, err := cluster.WaitForNActiveUnits(m, 0)
+		if err != nil {
+			fmt.Errorf("%v", err)
+		}
+		_, found := units["hello.service"]
+		if len(units) != 0 || !found {
+			fmt.Errorf("Expected hello.service to be sole active unit, got %v", units)
+		}
+	}
+
+	return nil
+}
+
 // destroyUnitRetrying() destroys the unit and ensure it disappears from the
 // unit list. It could take a little time until the unit gets destroyed.
 func destroyUnitRetrying(cluster platform.Cluster, m platform.Member, serviceFile string) error {