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() timeout := time.After(60 * time.Second) ticker := time.NewTicker(100 * time.Millisecond) 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) } 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 }