| 466 | } |
| 467 | |
| 468 | func waitAgentAPIStable(ctx context.Context, t testing.TB, apiClient *agentapisdk.Client, waitFor time.Duration, msg string) error { |
| 469 | t.Helper() |
| 470 | waitCtx, waitCancel := context.WithTimeout(ctx, waitFor) |
| 471 | defer waitCancel() |
| 472 | |
| 473 | start := time.Now() |
| 474 | var currStatus agentapisdk.AgentStatus |
| 475 | defer func() { |
| 476 | elapsed := time.Since(start) |
| 477 | t.Logf("%s: agent API status: %s (elapsed: %s)", msg, currStatus, elapsed.Round(100*time.Millisecond)) |
| 478 | }() |
| 479 | evts, errs, err := apiClient.SubscribeEvents(ctx) |
| 480 | require.NoError(t, err, "failed to subscribe to events") |
| 481 | for { |
| 482 | select { |
| 483 | case <-waitCtx.Done(): |
| 484 | return waitCtx.Err() |
| 485 | case evt := <-evts: |
| 486 | if esc, ok := evt.(agentapisdk.EventStatusChange); ok { |
| 487 | currStatus = esc.Status |
| 488 | if currStatus == agentapisdk.StatusStable { |
| 489 | return nil |
| 490 | } |
| 491 | } else { |
| 492 | var sb strings.Builder |
| 493 | if err := json.NewEncoder(&sb).Encode(evt); err != nil { |
| 494 | t.Logf("Failed to encode event: %v", err) |
| 495 | } |
| 496 | t.Logf("Got event: %s", sb.String()) |
| 497 | } |
| 498 | case err := <-errs: |
| 499 | return fmt.Errorf("read events: %w", err) |
| 500 | } |
| 501 | } |
| 502 | } |
| 503 | |
| 504 | // waitForMessagesWithCount retries GetMessages until it returns the expected number of messages or the timeout is reached. |
| 505 | func waitForMessagesWithCount(ctx context.Context, t testing.TB, apiClient *agentapisdk.Client, expectedCount int, timeout time.Duration, msg string) (*agentapisdk.GetMessagesResponse, error) { |