query-orchestration/internal/test/queryAPI/service.go

package queryapitest

import (
	"bytes"
	"context"
	"fmt"
	"log/slog"
	"mime/multipart"
	"net/http"
	"testing"
	"time"

	"queryorchestration/internal/test"

	queryapi "queryorchestration/pkg/queryAPI"

	"github.com/google/uuid"
	"github.com/oapi-codegen/runtime/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

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 SetQueryForClient(t testing.TB, client queryapi.ClientWithResponsesInterface, clientId queryapi.ClientID, queryId uuid.UUID) {
	t.Helper()

	newActiveVersion := int32(1)
	_, err := client.SetCollectorByClientIdWithResponse(t.Context(), clientId, queryapi.CollectorSet{
		ActiveVersion: &newActiveVersion,
		Fields: &[]queryapi.CollectorField{
			{
				Name:    "JSON_QUERY",
				QueryId: queryId,
			},
		},
	})
	require.NoError(t, err)
}

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)
		}
	}
}

func CreateDependentQueries(t testing.TB, client *queryapi.ClientWithResponses) (uuid.UUID, uuid.UUID) {
	contextQueryRes, err := client.CreateQueryWithResponse(t.Context(), queryapi.QueryCreate{
		Type: queryapi.CONTEXTFULL,
	})
	require.NoError(t, err)
	test.AssertStatus(t, http.StatusCreated, contextQueryRes.HTTPResponse)
	jcfg := `{"path":"keyone"}`
	jsonQueryRes, err := client.CreateQueryWithResponse(t.Context(), queryapi.QueryCreate{
		Type:            queryapi.JSONEXTRACTOR,
		Config:          &jcfg,
		RequiredQueries: &[]types.UUID{contextQueryRes.JSON201.Id},
	})
	require.NoError(t, err)
	test.AssertStatus(t, http.StatusCreated, jsonQueryRes.HTTPResponse)

	return contextQueryRes.JSON201.Id, jsonQueryRes.JSON201.Id
}