query-orchestration/api/queryAPI/query_test.go

package queryapi_test

import (
	"fmt"
	"net/http"
	"regexp"
	"testing"
	"time"

	"queryorchestration/internal/backgroundtask"
	"queryorchestration/internal/client"
	"queryorchestration/internal/collector"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document"
	"queryorchestration/internal/export"
	"queryorchestration/internal/query"
	"queryorchestration/internal/query/result"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue"
	"queryorchestration/internal/serviceconfig/queue/clientsync"
	"queryorchestration/internal/serviceconfig/queue/queryversionsync"
	"queryorchestration/internal/test"

	queryapi "queryorchestration/api/queryAPI"

	clientupdate "queryorchestration/internal/client/update"

	collectorset "queryorchestration/internal/collector/set"

	documentbatch "queryorchestration/internal/document/batch"
	documentupload "queryorchestration/internal/document/upload"

	resultprocessor "queryorchestration/internal/query/result/processor"
	querytest "queryorchestration/internal/query/test"
	queryupdate "queryorchestration/internal/query/update"

	"github.com/jackc/pgx/v5/pgtype"
	"github.com/stretchr/testify/require"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
)

func TestCreateQuery(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	body := queryapi.QueryCreate{
		Type: queryapi.CONTEXTFULL,
	}
	ctx, rec := createContextWithJSONBody(t, body)

	err := cons.CreateQuery(ctx)
	require.NoError(t, err)
	assert.Equal(t, http.StatusCreated, rec.Code)
	res := getBody[queryapi.IdMessage](t, rec)
	assert.NotEqual(t, uuid.Nil, res.Id)
}

func TestListQueries(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	ctx, rec := createContext(t)

	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeContextFull)
	require.NoError(t, err)

	err = cons.ListQueries(ctx)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.ListQueries{
		Queries: []queryapi.Query{
			{
				Id:            id,
				Type:          queryapi.CONTEXTFULL,
				ActiveVersion: 0,
				LatestVersion: 0,
			},
		},
	})
}

func TestGetQuery(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	ctx, rec := createContext(t)

	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeContextFull)
	require.NoError(t, err)

	err = cons.GetQuery(ctx, id)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.Query{
		Id:            id,
		Type:          queryapi.CONTEXTFULL,
		ActiveVersion: 0,
		LatestVersion: 0,
	})
}

func TestUpdateQuery(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.QueryVersionSyncURL = test.CreateQueue(t, cfg, test.QueryVersionSyncRunnerName)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeJsonExtractor)
	require.NoError(t, err)

	av := int32(1)
	c := `{"path": "val"}`
	body := queryapi.QueryUpdate{
		ActiveVersion: &av,
		Config:        &c,
	}
	ctx, rec := createContextWithJSONBody(t, body)

	err = cons.UpdateQuery(ctx, id)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assert.Empty(t, rec.Body.String())

	test.AssertMessageBody(t, cfg, cfg.GetQueryVersionSyncURL(), regexp.MustCompile(fmt.Sprintf(`{"id":"%s"}`, id)))
}

type ControllerConfig struct {
	serviceconfig.BaseConfig
	queryversionsync.QueryVersionSyncConfig
	clientsync.ClientSyncConfig
	queue.QueueConfig
	objectstore.ObjectStoreConfig
}

func (c *ControllerConfig) GetBackgroundRunner() *backgroundtask.Runner {
	return nil // Test config returns nil
}

func createControllerServices(cfg *ControllerConfig) *queryapi.Services {
	docsvc := document.New(cfg)
	col := collector.New(cfg)
	que := query.New(cfg)
	cli := client.New(cfg)
	return &queryapi.Services{
		Export:    export.New(),
		Collector: col,
		CollectorSet: collectorset.New(cfg, &collectorset.Services{
			Collector: col,
		}),
		Query: que,
		QueryUpdate: queryupdate.New(cfg, &queryupdate.Services{
			Query: que,
		}),
		QueryTest: querytest.New(cfg, &querytest.Services{
			Result: result.New(cfg, &result.Services{
				Query: que,
			}),
			Collector: col,
			Document:  docsvc,
		}),
		Client: cli,
		ClientUpdate: clientupdate.New(cfg, &clientupdate.Services{
			Client: cli,
		}),
		Document:       docsvc,
		DocumentUpload: documentupload.New(cfg),
		DocumentBatch:  documentbatch.New(cfg),
	}
}

func createDocumentWithText(t testing.TB, cfg *ControllerConfig, clientId string) (uuid.UUID, uuid.UUID, uuid.UUID) {
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: clientId,
		Hash:     "hash",
	})
	require.NoError(t, err)
	fill := "fill"
	cleanId, err := cfg.GetDBQueries().AddDocumentClean(t.Context(), &repository.AddDocumentCleanParams{
		Documentid: docId,
		Bucket:     &fill,
		Key:        &fill,
		Hash:       &fill,
		Mimetype: repository.NullCleanmimetype{
			Valid:         true,
			Cleanmimetype: repository.CleanmimetypeApplicationPdf,
		},
	})
	require.NoError(t, err)
	err = cfg.GetDBQueries().AddDocumentCleanEntry(t.Context(), &repository.AddDocumentCleanEntryParams{
		Cleanid: cleanId,
		Version: 1,
	})
	require.NoError(t, err)
	textId, err := cfg.GetDBQueries().AddDocumentText(t.Context(), &repository.AddDocumentTextParams{
		Cleanid: cleanId,
		Bucket:  fill,
		Key:     fill,
		Hash:    fill,
		Createdat: pgtype.Timestamp{
			Time:  time.Now(),
			Valid: true,
		},
	})
	require.NoError(t, err)
	err = cfg.GetDBQueries().AddDocumentTextEntry(t.Context(), &repository.AddDocumentTextEntryParams{
		Textid:  textId,
		Version: 1,
	})
	require.NoError(t, err)

	return docId, cleanId, textId
}

func TestTestQuery(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	contextQueryId, err := svc.Query.Create(t.Context(), &resultprocessor.Create{
		Type: resultprocessor.TypeContextFull,
	})
	require.NoError(t, err)

	_, err = svc.Client.Create(t.Context(), client.CreateParams{
		ID:   "client_id",
		Name: "client_name",
	})
	require.NoError(t, err)

	docId, _, textId := createDocumentWithText(t, cfg, "client_id")

	strVal := `{"mykey": "example_value", "oldkey": "old_value"}`
	_, err = cfg.GetDBQueries().AddResult(t.Context(), &repository.AddResultParams{
		Queryid:      contextQueryId,
		Value:        strVal,
		Textentryid:  textId,
		Queryversion: 1,
	})
	require.NoError(t, err)

	c := `{"path": "oldkey"}`
	queryId, err := svc.Query.Create(t.Context(), &resultprocessor.Create{
		Type:   resultprocessor.TypeJsonExtractor,
		Config: &c,
		RequiredQueryIDs: &[]uuid.UUID{
			contextQueryId,
		},
	})
	require.NoError(t, err)
	c = `{"path": "mykey"}`
	err = svc.QueryUpdate.Update(t.Context(), &resultprocessor.Update{
		ID:     queryId,
		Config: &c,
	})
	require.NoError(t, err)

	body := queryapi.QueryTestRequest{
		DocumentId:   docId,
		QueryVersion: 1,
	}
	ctx, rec := createContextWithJSONBody(t, body)

	err = cons.TestQuery(ctx, queryId)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.QueryTestResponse{
		Value: "old_value",
	})

	body = queryapi.QueryTestRequest{
		DocumentId:   docId,
		QueryVersion: 2,
	}
	ctx, rec = createContextWithJSONBody(t, body)

	err = cons.TestQuery(ctx, queryId)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.QueryTestResponse{
		Value: "example_value",
	})
}