Add pagination support for teams and users in org role tests. Enhanced mock GitHub server to handle paginated responses for teams and users, and updated tests to verify correct handling of grants and pagination scenarios.

Author: afalahi

pkg/connector/org_role_test.go

@@ -22,6 +22,10 @@ func TestOrgRole(t *testing.T) {
 
 		githubOrganization, _, _, githubUser, _ := mgh.Seed()
 
+		// Add user to org role
+		roleId := int64(1)
+		mgh.AddUserToOrgRole(roleId, *githubUser.ID)
+
 		githubClient := github.NewClient(mgh.Server())
 		cache := newOrgNameCache(githubClient)
 		client := orgRoleBuilder(githubClient, cache)
@@ -39,16 +43,19 @@ func TestOrgRole(t *testing.T) {
 			Resource: role,
 		}
 
+		// Grant the role to the user
 		grantAnnotations, err := client.Grant(ctx, user, &entitlement)
 		require.Nil(t, err)
 		require.Empty(t, grantAnnotations)
 
+		// Check that we can see both teams and users in the grants list
 		grants, nextToken, grantsAnnotations, err := client.Grants(ctx, role, &pagination.Token{})
 		require.Nil(t, err)
 		test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
-		require.Equal(t, "", nextToken)
-		require.Len(t, grants, 1)
+		require.Empty(t, nextToken) // No next token since we don't have a full page
+		require.Len(t, grants, 2)   // Should get both the team and user
 
+		// Revoke the role from the user
 		grant := v2.Grant{
 			Entitlement: &entitlement,
 			Principal:   user,
@@ -91,4 +98,41 @@ func TestOrgRole(t *testing.T) {
 		require.Empty(t, nextToken)
 		test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
 	})
+
+	t.Run("should handle pagination for teams and users", func(t *testing.T) {
+		mockGithub := mocks.NewMockGitHub()
+		githubOrganization, _, _, githubUser, _ := mockGithub.Seed()
+
+		// Add more teams to trigger pagination
+		for i := 0; i < 3; i++ {
+			teamId := int64(100 + i)
+			team := github.Team{
+				ID:           &teamId,
+				Organization: githubOrganization,
+			}
+			mockGithub.AddTeam(team)
+		}
+
+		// Add user to org role
+		roleId := int64(1)
+		mockGithub.AddUserToOrgRole(roleId, *githubUser.ID)
+
+		githubClient := github.NewClient(mockGithub.Server())
+		cache := newOrgNameCache(githubClient)
+		client := orgRoleBuilder(githubClient, cache)
+
+		organization, _ := organizationResource(ctx, githubOrganization, nil)
+		role, _ := orgRoleResource(ctx, &OrganizationRole{
+			ID:          1,
+			Name:        "Test Role",
+			Description: "Test Role Description",
+		}, organization)
+
+		// Test first page (should get all teams and users)
+		grants, nextToken, annotations, err := client.Grants(ctx, role, &pagination.Token{Size: 5})
+		require.Nil(t, err)
+		require.Empty(t, nextToken) // No next token since we got all results
+		require.Len(t, grants, 5)   // Should get all 4 teams (3 added + 1 from Seed) and 1 user
+		test.AssertNoRatelimitAnnotations(t, annotations)
+	})
 }

test/mocks/github.go

@@ -545,8 +545,36 @@ func (mgh MockGitHub) getOrgRoleTeams(
 		return
 	}
 
-	// Return empty teams list for now
-	_, _ = w.Write(mock.MustMarshal([]*github.Team{}))
+	// Get pagination parameters
+	page := 1
+	perPage := 30
+	if pageStr, ok := variables["page"]; ok {
+		if p, err := strconv.Atoi(pageStr); err == nil {
+			page = p
+		}
+	}
+	if perPageStr, ok := variables["per_page"]; ok {
+		if pp, err := strconv.Atoi(perPageStr); err == nil {
+			perPage = pp
+		}
+	}
+
+	// Return paginated teams
+	teams := make([]*github.Team, 0)
+	start := (page - 1) * perPage
+	end := start + perPage
+	i := 0
+	for _, team := range mgh.teams {
+		if i >= start && i < end {
+			teams = append(teams, &team)
+		}
+		i++
+		if i >= end {
+			break
+		}
+	}
+
+	_, _ = w.Write(mock.MustMarshal(teams))
 }
 
 func (mgh MockGitHub) getOrgRoleUsers(
@@ -560,10 +588,34 @@ func (mgh MockGitHub) getOrgRoleUsers(
 		return
 	}
 
+	// Get pagination parameters
+	page := 1
+	perPage := 30
+	if pageStr, ok := variables["page"]; ok {
+		if p, err := strconv.Atoi(pageStr); err == nil {
+			page = p
+		}
+	}
+	if perPageStr, ok := variables["per_page"]; ok {
+		if pp, err := strconv.Atoi(perPageStr); err == nil {
+			perPage = pp
+		}
+	}
+
+	// Return paginated users
 	users := make([]github.User, 0)
+	start := (page - 1) * perPage
+	end := start + perPage
+	i := 0
 	for _, userID := range memberships.ToSlice() {
-		if user, ok := mgh.users[userID]; ok {
-			users = append(users, user)
+		if i >= start && i < end {
+			if user, ok := mgh.users[userID]; ok {
+				users = append(users, user)
+			}
+		}
+		i++
+		if i >= end {
+			break
 		}
 	}
 
@@ -655,3 +707,24 @@ func (mgh MockGitHub) Server() *http.Client {
 	}
 	return mock.NewMockedHTTPClient(options...)
 }
+
+// AddTeam adds a team to the mock server for testing purposes.
+func (mgh *MockGitHub) AddTeam(team github.Team) {
+	mgh.teams[*team.ID] = team
+}
+
+// AddUserToOrgRole adds a user to an org role for testing purposes.
+func (mgh *MockGitHub) AddUserToOrgRole(roleID int64, userID int64) {
+	if _, ok := mgh.orgRoles[roleID]; !ok {
+		mgh.orgRoles[roleID] = mapset.NewSet[int64]()
+	}
+	mgh.orgRoles[roleID].Add(userID)
+}
+
+// AddTeamToOrgRole adds a team to an org role for testing purposes.
+func (mgh *MockGitHub) AddTeamToOrgRole(roleID int64, teamID int64) {
+	if _, ok := mgh.orgRoles[roleID]; !ok {
+		mgh.orgRoles[roleID] = mapset.NewSet[int64]()
+	}
+	mgh.orgRoles[roleID].Add(teamID)
+}