Accept newer revisions when waiting for policies to be assigned (#1289)

When tearing down tests, we try to reassign the previous policy. If it
was modified and we also expect the revision to be the same, it will
never succeed, it will wait till timeout.

internal/kibana/agents.go

@@ -16,7 +16,10 @@ import (
 	"github.com/elastic/elastic-package/internal/signal"
 )
 
-var waitForPolicyAssignedTimeout = 10 * time.Minute
+var (
+	waitForPolicyAssignedTimeout     = 10 * time.Minute
+	waitForPolicyAssignedRetryPeriod = 2 * time.Second
+)
 
 // Agent represents an Elastic Agent enrolled with fleet.
 type Agent struct {
@@ -83,12 +86,12 @@ func (c *Client) AssignPolicyToAgent(a Agent, p Policy) error {
 }
 
 func (c *Client) waitUntilPolicyAssigned(a Agent, p Policy) error {
-	timeout := time.Now().Add(waitForPolicyAssignedTimeout)
-	for {
-		if time.Now().After(timeout) {
-			return errors.New("timeout: policy hasn't been assigned in time")
-		}
+	timeout := time.NewTimer(waitForPolicyAssignedTimeout)
+	defer timeout.Stop()
+	ticker := time.NewTicker(waitForPolicyAssignedRetryPeriod)
+	defer ticker.Stop()
 
+	for {
 		if signal.SIGINT() {
 			return errors.New("SIGINT: cancel waiting for policy assigned")
 		}
@@ -99,13 +102,19 @@ func (c *Client) waitUntilPolicyAssigned(a Agent, p Policy) error {
 		}
 		logger.Debugf("Agent data: %s", agent.String())
 
-		if agent.PolicyID == p.ID && agent.PolicyRevision == p.Revision {
+		if agent.PolicyID == p.ID && agent.PolicyRevision >= p.Revision {
 			logger.Debugf("Policy revision assigned to the agent (ID: %s)...", a.ID)
 			break
 		}
 
 		logger.Debugf("Wait until the policy (ID: %s, revision: %d) is assigned to the agent (ID: %s)...", p.ID, p.Revision, a.ID)
-		time.Sleep(2 * time.Second)
+		select {
+		case <-timeout.C:
+			return errors.New("timeout: policy hasn't been assigned in time")
+		case <-ticker.C:
+			continue
+		}
+
 	}
 	return nil
 }