package queryapitest import ( "bytes" "context" "fmt" "log/slog" "mime/multipart" "net/http" "testing" "time" "queryorchestration/internal/test" queryapi "queryorchestration/pkg/queryAPI" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // Note: Query-related helper functions have been removed. // See remove_query_plan.md for details. type File struct { ClientID queryapi.ClientID Filename string Content []byte } func CreateFile(t testing.TB, client queryapi.ClientWithResponsesInterface, file File) { var buf bytes.Buffer fileBody := multipart.NewWriter(&buf) name := file.Filename if name == "" { name = "helloworld" } write, err := fileBody.CreateFormFile("file", name) require.NoError(t, err) _, err = write.Write(file.Content) require.NoError(t, err) err = fileBody.Close() require.NoError(t, err) body := bytes.NewReader(buf.Bytes()) uploadRes, err := client.UploadDocumentWithBodyWithResponse(t.Context(), file.ClientID, fileBody.FormDataContentType(), body) require.NoError(t, err) require.Equal(t, http.StatusOK, uploadRes.StatusCode()) } func CreateClientWithSync(t testing.TB, client queryapi.ClientWithResponsesInterface) queryapi.ClientID { t.Helper() clientCreateRes, err := client.CreateClientWithResponse(t.Context(), queryapi.ClientCreate{ Name: "example_name", Id: "ID", }) require.NoError(t, err) require.NotNil(t, clientCreateRes.JSON201, "expected 201 response but got status %d", clientCreateRes.StatusCode()) require.NotEmpty(t, clientCreateRes.JSON201.Id, "client ID should not be empty") canSync := true _, err = client.UpdateClientWithResponse(t.Context(), clientCreateRes.JSON201.Id, queryapi.ClientUpdate{ CanSync: &canSync, }) require.NoError(t, err) return clientCreateRes.JSON201.Id } func WaitForClientStatus(t testing.TB, ctx context.Context, service queryapi.ClientWithResponsesInterface, id string, status queryapi.ClientStatus) { t.Helper() // Use configurable polling timeout (longer in CI environments to handle slower infrastructure) // and configurable poll interval to balance responsiveness with rate limiting concerns timeout := time.After(test.GetPollingTimeout()) ticker := time.NewTicker(test.GetPollInterval()) defer ticker.Stop() for { select { case <-timeout: require.NoError(t, fmt.Errorf("Timeout waiting for client status to become %s", status)) case <-ticker.C: jRes, err := service.GetStatusByClientIdWithResponse(ctx, id) if err != nil { slog.Error("error getting status", "error", err.Error()) require.NoError(t, err) } // Check if we got a non-200 response (e.g., 429 rate limit) if jRes.JSON200 == nil { if jRes.StatusCode() == http.StatusTooManyRequests { slog.Warn("rate limit hit while polling status, will retry", "status_code", jRes.StatusCode()) continue } require.Fail(t, fmt.Sprintf("unexpected status code %d waiting for client status", jRes.StatusCode())) return } if jRes.JSON200.Status == status { assert.Equal(t, status, jRes.JSON200.Status) slog.Info("returned client status", "status", jRes.JSON200.Status) return } slog.Error("unexpected status", "status", jRes.JSON200.Status) } } }