Initial QSfera import

Server/services/search/pkg/opensearch/internal/test/os.go

@@ -0,0 +1,256 @@
+package opensearchtest
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"slices"
+	"testing"
+
+	opensearchgo "github.com/opensearch-project/opensearch-go/v4"
+	opensearchgoAPI "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
+	"github.com/stretchr/testify/require"
+
+	"github.com/qsfera/server/services/search/pkg/config"
+)
+
+type TestClient struct {
+	c       *opensearchgoAPI.Client
+	Require *testRequireClient
+}
+
+func NewDefaultTestClient(t *testing.T, cfg config.EngineOpenSearchClient) *TestClient {
+	client, err := opensearchgoAPI.NewClient(opensearchgoAPI.Config{
+		Client: opensearchgo.Config{
+			Addresses: cfg.Addresses,
+			Username:  cfg.Username,
+			Password:  cfg.Password,
+		},
+	})
+	require.NoError(t, err, "failed to create OpenSearch client")
+
+	return NewTestClient(t, client)
+}
+
+func NewTestClient(t *testing.T, client *opensearchgoAPI.Client) *TestClient {
+	tc := &TestClient{c: client}
+	trc := &testRequireClient{tc: tc, t: t}
+	tc.Require = trc
+
+	return tc
+}
+
+func (tc *TestClient) Client() *opensearchgoAPI.Client {
+	return tc.c
+}
+
+func (tc *TestClient) IndicesReset(ctx context.Context, indices []string) error {
+	indicesToDelete := make([]string, 0, len(indices))
+	for _, index := range indices {
+		exist, err := tc.IndicesExists(ctx, []string{index})
+		if err != nil {
+			return fmt.Errorf("failed to check if index %s exists: %w", index, err)
+		}
+
+		if !exist {
+			continue
+		}
+
+		indicesToDelete = append(indicesToDelete, index)
+	}
+
+	if len(indicesToDelete) == 0 {
+		// If no indices to delete, return nil
+		return nil
+	}
+
+	return tc.IndicesDelete(ctx, indicesToDelete)
+}
+
+func (tc *TestClient) IndicesExists(ctx context.Context, indices []string) (bool, error) {
+	if err := tc.IndicesRefresh(ctx, indices, []int{404}); err != nil {
+		return false, err
+	}
+
+	resp, err := tc.c.Indices.Exists(ctx, opensearchgoAPI.IndicesExistsReq{
+		Indices: indices,
+	})
+	switch {
+	case resp != nil && resp.StatusCode == 404:
+		return false, nil
+	case err != nil:
+		return false, fmt.Errorf("failed to check if indices exist: %w", err)
+	case resp != nil && resp.IsError():
+		return false, fmt.Errorf("failed to check if indices exist: %s", resp.String())
+	default:
+		return true, nil
+	}
+}
+
+func (tc *TestClient) IndicesRefresh(ctx context.Context, indices []string, allow []int) error {
+	resp, err := tc.c.Indices.Refresh(ctx, &opensearchgoAPI.IndicesRefreshReq{
+		Indices: indices,
+	})
+
+	isAllowed := resp != nil
+	isAllowed = isAllowed && resp.Inspect().Response != nil
+	isAllowed = isAllowed && slices.Contains(allow, resp.Inspect().Response.StatusCode)
+
+	if err != nil && !isAllowed {
+		return fmt.Errorf("failed to refresh indices %v: %w", indices, err)
+	}
+
+	return nil
+}
+
+func (tc *TestClient) IndicesDelete(ctx context.Context, indices []string) error {
+	if err := tc.IndicesRefresh(ctx, indices, []int{}); err != nil {
+		return err
+	}
+
+	resp, err := tc.c.Indices.Delete(ctx, opensearchgoAPI.IndicesDeleteReq{
+		Indices: indices,
+	})
+	switch {
+	case err != nil:
+		return fmt.Errorf("failed to delete indices: %w", err)
+	case !resp.Acknowledged:
+		return errors.New("indices deletion not acknowledged")
+	default:
+		return nil
+	}
+}
+
+func (tc *TestClient) IndicesCreate(ctx context.Context, index string, body io.Reader) error {
+	resp, err := tc.c.Indices.Create(ctx, opensearchgoAPI.IndicesCreateReq{
+		Index: index,
+		Body:  body,
+	})
+
+	switch {
+	case err != nil:
+		return fmt.Errorf("failed to create index %s: %w", index, err)
+	case !resp.Acknowledged:
+		return fmt.Errorf("index creation not acknowledged for %s", index)
+	default:
+		return nil
+	}
+}
+
+func (tc *TestClient) IndicesCount(ctx context.Context, indices []string, body io.Reader) (int, error) {
+	if err := tc.IndicesRefresh(ctx, indices, []int{404}); err != nil {
+		return 0, err
+	}
+
+	resp, err := tc.c.Indices.Count(ctx, &opensearchgoAPI.IndicesCountReq{
+		Indices: indices,
+		Body:    body,
+	})
+
+	switch {
+	case err != nil:
+		return 0, fmt.Errorf("failed to count documents in indices: %w", err)
+	default:
+		return resp.Count, nil
+	}
+}
+
+func (tc *TestClient) DocumentCreate(ctx context.Context, index, id string, body io.Reader) error {
+	if err := tc.IndicesRefresh(ctx, []string{index}, []int{404}); err != nil {
+		return err
+	}
+
+	_, err := tc.c.Document.Create(ctx, opensearchgoAPI.DocumentCreateReq{
+		Index:      index,
+		DocumentID: id,
+		Body:       body,
+	})
+	switch {
+	case err != nil:
+		return fmt.Errorf("failed to create document in index %s: %w", index, err)
+	default:
+		return nil
+	}
+}
+
+func (tc *TestClient) Update(ctx context.Context, index, id string, body io.Reader) error {
+	if err := tc.IndicesRefresh(ctx, []string{index}, []int{404}); err != nil {
+		return err
+	}
+
+	_, err := tc.c.Update(ctx, opensearchgoAPI.UpdateReq{
+		Index:      index,
+		DocumentID: id,
+		Body:       body,
+	})
+	switch {
+	case err != nil:
+		return fmt.Errorf("failed to update document in index %s: %w", index, err)
+	default:
+		return nil
+	}
+}
+
+func (tc *TestClient) Search(ctx context.Context, index string, body io.Reader) (opensearchgoAPI.SearchHits, error) {
+	if err := tc.IndicesRefresh(ctx, []string{index}, []int{404}); err != nil {
+		return opensearchgoAPI.SearchHits{}, err
+	}
+
+	resp, err := tc.c.Search(ctx, &opensearchgoAPI.SearchReq{
+		Indices: []string{index},
+		Body:    body,
+	})
+	if err != nil {
+		return opensearchgoAPI.SearchHits{}, fmt.Errorf("failed to search in index %s: %w", index, err)
+	}
+
+	return resp.Hits, nil
+}
+
+type testRequireClient struct {
+	tc *TestClient
+	t  *testing.T
+}
+
+func (trc *testRequireClient) IndicesReset(indices []string) {
+	require.NoError(trc.t, trc.tc.IndicesReset(trc.t.Context(), indices))
+}
+
+func (trc *testRequireClient) IndicesRefresh(indices []string, ignore []int) {
+	require.NoError(trc.t, trc.tc.IndicesRefresh(trc.t.Context(), indices, ignore))
+}
+
+func (trc *testRequireClient) IndicesCreate(index string, body io.Reader) {
+	require.NoError(trc.t, trc.tc.IndicesCreate(trc.t.Context(), index, body))
+}
+
+func (trc *testRequireClient) IndicesDelete(indices []string) {
+	require.NoError(trc.t, trc.tc.IndicesDelete(trc.t.Context(), indices))
+}
+
+func (trc *testRequireClient) IndicesCount(indices []string, body io.Reader, expected int) {
+	count, err := trc.tc.IndicesCount(trc.t.Context(), indices, body)
+
+	switch {
+	case expected <= 0:
+		require.True(trc.t, count <= 0, "expected indices to have no documents, but got a count of %d", count)
+	default:
+		require.Equal(trc.t, expected, count, "expected indices to have %d documents, but got %d", expected, count)
+		require.NoError(trc.t, err, "expected indices to have documents, but got an error")
+	}
+}
+
+func (trc *testRequireClient) DocumentCreate(index, id string, body io.Reader) {
+	require.NoError(trc.t, trc.tc.DocumentCreate(trc.t.Context(), index, id, body))
+}
+
+func (trc *testRequireClient) Update(index, id string, body io.Reader) {
+	require.NoError(trc.t, trc.tc.Update(trc.t.Context(), index, id, body))
+}
+
+func (trc *testRequireClient) Search(index string, body io.Reader) opensearchgoAPI.SearchHits {
+	hits, err := trc.tc.Search(trc.t.Context(), index, body)
+	require.NoError(trc.t, err)
+	return hits
+}