Add pagination when getting user groups (#615)

* Add pagination when getting user groups

* Support paginated requests for setGroups

* Cleanup tests for reusability

Co-authored-by: Brent Souza <bsouza@acadian-asset.com>

okta/user.go

@@ -290,33 +290,63 @@ func setAdminRoles(ctx context.Context, d *schema.ResourceData, m interface{}) e
 
 // set all groups currently attached to the user
 func setAllGroups(ctx context.Context, d *schema.ResourceData, c *okta.Client) error {
-	groups, _, err := c.User.ListUserGroups(ctx, d.Id())
+	groups, response, err := c.User.ListUserGroups(ctx, d.Id())
 	if err != nil {
 		return fmt.Errorf("failed to list user groups: %v", err)
 	}
-	groupIDs := make([]interface{}, len(groups))
-	for i := range groups {
-		groupIDs[i] = groups[i].Id
+
+	groupIDs := make([]interface{}, 0)
+
+	for {
+		for _, group := range groups {
+			groupIDs = append(groupIDs, group.Id)
+		}
+
+		if !response.HasNextPage() {
+			break
+		}
+
+		response, err = response.Next(ctx, &groups)
+
+		if err != nil {
+			return fmt.Errorf("failed to list user groups: %v", err)
+		}
 	}
+
 	return setNonPrimitives(d, map[string]interface{}{
 		"group_memberships": schema.NewSet(schema.HashString, groupIDs),
 	})
 }
 
 // set groups attached to the user that can be changed
 func setGroups(ctx context.Context, d *schema.ResourceData, c *okta.Client) error {
-	groups, _, err := c.User.ListUserGroups(ctx, d.Id())
+	groups, response, err := c.User.ListUserGroups(ctx, d.Id())
 	if err != nil {
 		return fmt.Errorf("failed to list user groups: %v", err)
 	}
+
 	groupIDs := make([]interface{}, 0)
-	// ignore saving build-in or app groups into state so we don't end up with perpetual diffs,
-	// because it's impossible to remove user from build-in or app group via API
-	for _, group := range groups {
-		if group.Type != "BUILT_IN" && group.Type != "APP_GROUP" {
-			groupIDs = append(groupIDs, group.Id)
+
+	for {
+		// ignore saving build-in or app groups into state so we don't end up with perpetual diffs,
+		// because it's impossible to remove user from build-in or app group via API
+		for _, group := range groups {
+			if group.Type != "BUILT_IN" && group.Type != "APP_GROUP" {
+				groupIDs = append(groupIDs, group.Id)
+			}
+		}
+
+		if !response.HasNextPage() {
+			break
+		}
+
+		response, err = response.Next(ctx, &groups)
+
+		if err != nil {
+			return fmt.Errorf("failed to list user groups: %v", err)
 		}
 	}
+
 	return setNonPrimitives(d, map[string]interface{}{
 		"group_memberships": schema.NewSet(schema.HashString, groupIDs),
 	})

okta/user_test.go

@@ -0,0 +1,105 @@
+package okta
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"sort"
+	"testing"
+
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/okta/okta-sdk-golang/v2/okta"
+)
+
+func getGroupPagesJson(t *testing.T) ([]byte, []byte) {
+	firstPageOfGroups, err := json.Marshal([]*okta.Group{
+		{Id: "foo"},
+		{Id: "bar"},
+	})
+
+	if err != nil {
+		t.Fatalf("could not serialize first set of groups: %s", err)
+	}
+
+	secondPageOfGroups, err := json.Marshal([]*okta.Group{
+		{Id: "baz"},
+		{Id: "qux"},
+	})
+
+	if err != nil {
+		t.Fatalf("could not serialize second set of groups: %s", err)
+	}
+
+	return firstPageOfGroups, secondPageOfGroups
+}
+
+type userGroupFunc func(ctx context.Context, d *schema.ResourceData, c *okta.Client) error
+
+func testUserGroupFetchesAllPages(t *testing.T, fn userGroupFunc) {
+	s := dataSourceUser().Schema
+	d := schema.TestResourceDataRaw(t, s, map[string]interface{}{
+		"Id": "foo",
+	})
+
+	firstPageOfGroups, secondPageOfGroups := getGroupPagesJson(t)
+
+	ctx, c, err := newTestOktaClientWithResponse(func(req *http.Request) *http.Response {
+		q := req.URL.Query()
+
+		if q.Has("after") {
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewReader(secondPageOfGroups)),
+				Header:     make(http.Header),
+			}
+		}
+
+		return &http.Response{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewReader(firstPageOfGroups)),
+			Header: http.Header{
+				"Link": []string{"<https://foo.okta.com?limit=2&after=0>; rel=\"next\""},
+			},
+		}
+	})
+
+	if err != nil {
+		t.Fatalf("could not create an okta client instance: %s", err)
+	}
+
+	err = fn(ctx, d, c)
+
+	if err != nil {
+		t.Fatalf("fetching groups failed: %s", err)
+	}
+
+	groups := convertInterfaceToStringSetNullable(d.Get("group_memberships"))
+
+	if len(groups) != 4 {
+		t.Fatalf("expected 4 groups; got %d", len(groups))
+	}
+
+	expected := []string{"bar", "baz", "foo", "qux"}
+	sort.Strings(groups)
+	allMatch := true
+
+	for i, group := range expected {
+		if group != groups[i] {
+			allMatch = false
+		}
+	}
+
+	if !allMatch {
+		t.Fatalf("expected %s; got %s", expected, groups)
+	}
+}
+
+func TestUserSetAllGroups(t *testing.T) {
+	testUserGroupFetchesAllPages(t, setAllGroups)
+}
+
+func TestUserSetGroups(t *testing.T) {
+	testUserGroupFetchesAllPages(t, setGroups)
+}

okta/utils_for_test.go

@@ -1,11 +1,14 @@
 package okta
 
 import (
+	"context"
 	"fmt"
+	"net/http"
 	"strings"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/okta/okta-sdk-golang/v2/okta"
 )
 
 type checkUpstream func(string) (bool, error)
@@ -69,3 +72,34 @@ func condenseError(errorList []error) error {
 	}
 	return fmt.Errorf("series of errors occurred: %s", strings.Join(msgList, ", "))
 }
+
+type roundTripFunc func(req *http.Request) *http.Response
+
+func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req), nil
+}
+
+func newTestHttpClient(fn roundTripFunc) *http.Client {
+	return &http.Client{
+		Transport: fn,
+	}
+}
+
+func newTestOktaClientWithResponse(response roundTripFunc) (context.Context, *okta.Client, error) {
+	ctx := context.Background()
+
+	h := newTestHttpClient(response)
+
+	oktaCtx, c, err := okta.NewClient(
+		ctx,
+		okta.WithOrgUrl("https://foo.okta.com"),
+		okta.WithToken("f0oT0k3n"),
+		okta.WithHttpClientPtr(h),
+	)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return oktaCtx, c, nil
+}