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 }