From 4599af44a42209ae46091e2c8593308e497b42bc Mon Sep 17 00:00:00 2001
From: Tomer Heber <theber@zscaler.com>
Date: Tue, 19 Nov 2024 18:54:29 -0600
Subject: [PATCH] Fix: add pagination for teams api

---
 client/team.go      | 31 ++++++++++++++++++++++++-------
 client/team_test.go | 30 ++++++++++++++++++++++++++----
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/client/team.go b/client/team.go
index 64189843..4f11a2c5 100644
--- a/client/team.go
+++ b/client/team.go
@@ -22,6 +22,11 @@ type Team struct {
 	Users          []User `json:"users"`
 }
 
+type PaginatedTeamsResponse struct {
+	Teams       []Team `json:"teams"`
+	NextPageKey string `json:"nextPageKey"`
+}
+
 func (client *ApiClient) TeamCreate(payload TeamCreatePayload) (Team, error) {
 	if payload.Name == "" {
 		return Team{}, errors.New("must specify team name on creation")
@@ -84,20 +89,32 @@ func (client *ApiClient) GetTeams(params map[string]string) ([]Team, error) {
 		return nil, err
 	}
 
-	var result []Team
+	var teams []Team
 
-	err = client.http.Get("/teams/organizations/"+organizationId, params, &result)
-	if err != nil {
-		return nil, err
+	for {
+		var res PaginatedTeamsResponse
+
+		if err := client.http.Get("/teams/organizations/"+organizationId, params, &res); err != nil {
+			return nil, err
+		}
+
+		teams = append(teams, res.Teams...)
+
+		nextPageKey := res.NextPageKey
+		if nextPageKey == "" {
+			break
+		}
+
+		params["offset"] = nextPageKey
 	}
 
-	return result, err
+	return teams, nil
 }
 
 func (client *ApiClient) Teams() ([]Team, error) {
-	return client.GetTeams(nil)
+	return client.GetTeams(map[string]string{"limit": "100"})
 }
 
 func (client *ApiClient) TeamsByName(name string) ([]Team, error) {
-	return client.GetTeams(map[string]string{"name": name})
+	return client.GetTeams(map[string]string{"name": name, "limit": "100"})
 }
diff --git a/client/team_test.go b/client/team_test.go
index 7890a9be..703d5b5e 100644
--- a/client/team_test.go
+++ b/client/team_test.go
@@ -16,6 +16,11 @@ var _ = Describe("Teams Client", func() {
 		Name: "team-name",
 	}
 
+	mockTeam2 := Team{
+		Id:   "team-id2",
+		Name: "team-name2",
+	}
+
 	Describe("Get Single Team", func() {
 		var returnedTeam Team
 
@@ -40,13 +45,29 @@ var _ = Describe("Teams Client", func() {
 	Describe("Get All Teams", func() {
 		var returnedTeams []Team
 		mockTeams := []Team{mockTeam}
+		mockTeams2 := []Team{mockTeam2}
 
 		BeforeEach(func() {
 			mockOrganizationIdCall()
 			httpCall = mockHttpClient.EXPECT().
-				Get("/teams/organizations/"+organizationId, nil, gomock.Any()).
-				Do(func(path string, request interface{}, response *[]Team) {
-					*response = mockTeams
+				Get("/teams/organizations/"+organizationId, map[string]string{
+					"limit": "100",
+				}, gomock.Any()).
+				Do(func(path string, request interface{}, response *PaginatedTeamsResponse) {
+					*response = PaginatedTeamsResponse{
+						Teams:       mockTeams,
+						NextPageKey: "next_page_key",
+					}
+				})
+			httpCall2 = mockHttpClient.EXPECT().
+				Get("/teams/organizations/"+organizationId, map[string]string{
+					"offset": "next_page_key",
+					"limit":  "100",
+				}, gomock.Any()).
+				Do(func(path string, request interface{}, response *PaginatedTeamsResponse) {
+					*response = PaginatedTeamsResponse{
+						Teams: mockTeams2,
+					}
 				})
 			returnedTeams, _ = apiClient.Teams()
 		})
@@ -57,10 +78,11 @@ var _ = Describe("Teams Client", func() {
 
 		It("Should send GET request", func() {
 			httpCall.Times(1)
+			httpCall2.Times(1)
 		})
 
 		It("Should return teams", func() {
-			Expect(returnedTeams).To(Equal(mockTeams))
+			Expect(returnedTeams).To(Equal(append(mockTeams, mockTeams2...)))
 		})
 	})