Follow next_page_token in SearchEntitlements

pkg/client/entitlement.go

@@ -98,35 +98,61 @@ func (e *ExpandableEntitlementWithBindings) SetPath(pathname string, value int)
 }
 
 func (c *client) SearchEntitlements(ctx context.Context, filter *SearchEntitlementsFilter) ([]*EntitlementWithBindings, error) {
-	// TODO(morgabra) Pagination
 	// TODO(morgabra) Should we abstract the OpenAPI objects from the rest of cone? Kinda... no? But they aren't typed...
-	req := shared.RequestCatalogSearchServiceSearchEntitlementsRequest{
-		EntitlementAlias: stringPtr(filter.EntitlementAlias),
-		GrantedStatus:    filter.GrantedStatus.ToPointer(),
-		PageSize:         intPtr(100),
-		PageToken:        nil,
-		Query:            stringPtr(filter.Query),
-		AppDisplayName:   stringPtr(filter.AppDisplayName),
-		IncludeDeleted:   &filter.IncludeDeleted,
-		ExpandMask:       &filter.AppEntitlementExpandMask,
-	}
-	resp, err := c.sdk.RequestCatalogSearch.SearchEntitlements(ctx, &req)
-	if err != nil {
-		return nil, err
-	}
+	var rv []*EntitlementWithBindings
+	var pageToken *string
+	for page := 0; ; page++ {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 
-	if err := NewHTTPError(resp.RawResponse); err != nil {
-		return nil, err
-	}
+		req := shared.RequestCatalogSearchServiceSearchEntitlementsRequest{
+			EntitlementAlias: stringPtr(filter.EntitlementAlias),
+			GrantedStatus:    filter.GrantedStatus.ToPointer(),
+			PageSize:         intPtr(100),
+			PageToken:        pageToken,
+			Query:            stringPtr(filter.Query),
+			AppDisplayName:   stringPtr(filter.AppDisplayName),
+			IncludeDeleted:   &filter.IncludeDeleted,
+			ExpandMask:       &filter.AppEntitlementExpandMask,
+		}
+		resp, err := c.sdk.RequestCatalogSearch.SearchEntitlements(ctx, &req)
+		if err != nil {
+			return nil, err
+		}
 
-	list := resp.RequestCatalogSearchServiceSearchEntitlementsResponse.List
-	if list == nil {
-		return nil, errors.New("search-entitlements: list is nil")
+		if err := NewHTTPError(resp.RawResponse); err != nil {
+			return nil, err
+		}
+
+		body := resp.RequestCatalogSearchServiceSearchEntitlementsResponse
+		if body == nil || (body.List == nil && page == 0) {
+			return nil, errors.New("search-entitlements: list is nil")
+		}
+
+		// Expanded paths (e.g. "$.expanded[0]") are indexes into this response's own
+		// Expanded slice, so each page is resolved against its own expanded objects
+		// before the results are accumulated.
+		pageResults, err := convertSearchEntitlementsPage(body)
+		if err != nil {
+			return nil, err
+		}
+		rv = append(rv, pageResults...)
+
+		if body.NextPageToken == nil || *body.NextPageToken == "" {
+			break
+		}
+		pageToken = body.NextPageToken
 	}
+	return rv, nil
+}
+
+func convertSearchEntitlementsPage(body *shared.RequestCatalogSearchServiceSearchEntitlementsResponse) ([]*EntitlementWithBindings, error) {
+	list := body.List
 
 	// Unmarshal the expanded fields
-	expanded := make([]any, 0, len(resp.RequestCatalogSearchServiceSearchEntitlementsResponse.Expanded))
-	for _, x := range resp.RequestCatalogSearchServiceSearchEntitlementsResponse.Expanded {
+	expanded := make([]any, 0, len(body.Expanded))
+	for _, x := range body.Expanded {
 		x := x
 		converted, err := UnmarshalAnyType[shared.RequestCatalogSearchServiceSearchEntitlementsResponseExpanded](&x)
 		if err != nil {
@@ -147,10 +173,9 @@ func (c *client) SearchEntitlements(ctx context.Context, filter *SearchEntitleme
 	}
 
 	// Populate the expandable objects with the indexes of related objects
-	err = ExpandableReponse[*ExpandableEntitlementWithBindings]{
+	err := ExpandableReponse[*ExpandableEntitlementWithBindings]{
 		List: expandableList,
 	}.PopulateExpandedIndexes()
-
 	if err != nil {
 		return nil, err
 	}

pkg/client/entitlement_test.go

@@ -0,0 +1,102 @@
+package client
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/conductorone/conductorone-sdk-go/pkg/models/shared"
+)
+
+// searchPage builds a search response page whose single entitlement points at
+// its own expanded[0] app, mirroring the server's per-response indexing.
+func searchPage(entID, appID, nextToken string) map[string]any {
+	page := map[string]any{
+		"list": []any{
+			map[string]any{
+				"entitlement": map[string]any{
+					"appEntitlement": map[string]any{"id": entID, "appId": appID},
+					"appPath":        "$.expanded[0]",
+				},
+			},
+		},
+		"expanded": []any{
+			map[string]any{"@type": atTypeApp, "id": appID, "displayName": "App " + appID},
+		},
+	}
+	if nextToken != "" {
+		page["nextPageToken"] = nextToken
+	}
+	return page
+}
+
+func TestSearchEntitlementsFollowsPageToken(t *testing.T) {
+	var tokens []string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var req shared.RequestCatalogSearchServiceSearchEntitlementsRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			t.Errorf("Decode() unexpected error: %v", err)
+		}
+		if req.PageSize == nil || *req.PageSize != 100 {
+			t.Errorf("page size = %v, want 100", req.PageSize)
+		}
+		tokens = append(tokens, StringFromPtr(req.PageToken))
+
+		var page map[string]any
+		switch StringFromPtr(req.PageToken) {
+		case "":
+			page = searchPage("e1", "a1", "t1")
+		case "t1":
+			page = searchPage("e2", "a2", "t2")
+		case "t2":
+			page = searchPage("e3", "a3", "")
+		default:
+			t.Errorf("unexpected page token %q", StringFromPtr(req.PageToken))
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(page)
+	}))
+	defer server.Close()
+
+	c := newPaperSecretTestClient(server.URL, server.Client())
+	got, err := c.SearchEntitlements(context.Background(), &SearchEntitlementsFilter{})
+	if err != nil {
+		t.Fatalf("SearchEntitlements() unexpected error: %v", err)
+	}
+
+	if want := []string{"", "t1", "t2"}; fmt.Sprint(tokens) != fmt.Sprint(want) {
+		t.Fatalf("page tokens = %v, want %v", tokens, want)
+	}
+	if len(got) != 3 {
+		t.Fatalf("got %d entitlements, want 3", len(got))
+	}
+	for i, id := range []string{"1", "2", "3"} {
+		if gotID := StringFromPtr(got[i].Entitlement.ID); gotID != "e"+id {
+			t.Errorf("entitlement[%d] id = %q, want e%s", i, gotID, id)
+		}
+		// Every page uses expanded[0]; each entitlement must resolve to its own page's app.
+		app := GetExpanded[shared.App](got[i], ExpandedApp)
+		if app == nil || StringFromPtr(app.ID) != "a"+id {
+			t.Errorf("entitlement[%d] expanded app = %v, want a%s", i, app, id)
+		}
+	}
+}
+
+func TestSearchEntitlementsHonorsContextCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		// Always advertise another page; cancel after the first response.
+		cancel()
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(searchPage("e", "a", "next"))
+	}))
+	defer server.Close()
+
+	c := newPaperSecretTestClient(server.URL, server.Client())
+	if _, err := c.SearchEntitlements(ctx, &SearchEntitlementsFilter{}); err == nil {
+		t.Fatal("SearchEntitlements() error = nil, want context error")
+	}
+}