support org role tests by updating Seed method to include OrganizationRole. Refactor test cases to utilize the new role data

Author: afalahi

pkg/connector/org_role_test.go

@@ -20,42 +20,57 @@ func TestOrgRole(t *testing.T) {
 	t.Run("should grant and revoke entitlements", func(t *testing.T) {
 		mgh := mocks.NewMockGitHub()
 
-		githubOrganization, _, _, githubUser, _ := mgh.Seed()
-
-		// Add user to org role
-		roleId := int64(1)
-		mgh.AddUserToOrgRole(roleId, *githubUser.ID)
+		githubOrganization, _, _, githubUser, orgRole, _ := mgh.Seed()
 
 		githubClient := github.NewClient(mgh.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",
+		roleResource, _ := orgRoleResource(ctx, &OrganizationRole{
+			ID:          orgRole.ID,
+			Name:        orgRole.Name,
+			Description: orgRole.Description,
 		}, organization)
 		user, _ := userResource(ctx, githubUser, *githubUser.Email, nil)
 
 		entitlement := v2.Entitlement{
-			Id:       entitlement2.NewEntitlementID(role, "assigned"),
-			Resource: role,
+			Id:       entitlement2.NewEntitlementID(roleResource, "assigned"),
+			Resource: roleResource,
 		}
 
 		// 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.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
+		grants := make([]*v2.Grant, 0)
+		bag := &pagination.Bag{}
+		for {
+			pToken := pagination.Token{}
+			state := bag.Current()
+			if state != nil {
+				token, _ := bag.Marshal()
+				pToken.Token = token
+			}
+
+			nextGrants, nextToken, grantsAnnotations, err := client.Grants(ctx, roleResource, &pToken)
+			grants = append(grants, nextGrants...)
+
+			require.Nil(t, err)
+			test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+			if nextToken == "" {
+				break
+			}
+
+			err = bag.Unmarshal(nextToken)
+			if err != nil {
+				t.Error(err)
+			}
+		}
+
+		require.Len(t, grants, 2)
 
-		// Revoke the role from the user
 		grant := v2.Grant{
 			Entitlement: &entitlement,
 			Principal:   user,
@@ -70,7 +85,7 @@ func TestOrgRole(t *testing.T) {
 		mockGithub := mocks.NewMockGitHub()
 		mockGithub.SimulateOrgRolePermErr = true
 
-		githubOrganization, _, _, _, _ := mockGithub.Seed()
+		githubOrganization, _, _, _, orgRole, _ := mockGithub.Seed()
 
 		githubClient := github.NewClient(mockGithub.Server())
 		cache := newOrgNameCache(githubClient)
@@ -87,52 +102,16 @@ func TestOrgRole(t *testing.T) {
 
 		// Test Grants with permission error
 		role, _ := orgRoleResource(ctx, &OrganizationRole{
-			ID:          1,
-			Name:        "Test Role",
-			Description: "Test Role Description",
+			ID:          orgRole.ID,
+			Name:        orgRole.Name,
+			Description: orgRole.Description,
 		}, organization)
 
 		grants, nextToken, grantsAnnotations, err := client.Grants(ctx, role, &pagination.Token{})
 		require.Nil(t, err)
 		require.Empty(t, grants)
-		require.Empty(t, nextToken)
+		// The token should contain the initial state for users
+		require.NotEmpty(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)
-	})
 }

pkg/connector/org_test.go

@@ -19,7 +19,7 @@ func TestOrganization(t *testing.T) {
 	t.Run("should grant and revoke entitlements", func(t *testing.T) {
 		mgh := mocks.NewMockGitHub()
 
-		githubOrganization, _, _, githubUser, _ := mgh.Seed()
+		githubOrganization, _, _, githubUser, _, _ := mgh.Seed()
 
 		githubClient := github.NewClient(mgh.Server())
 		cache := newOrgNameCache(githubClient)

pkg/connector/repository_test.go

@@ -20,7 +20,7 @@ func TestRepository(t *testing.T) {
 	t.Run("should grant and revoke entitlements", func(t *testing.T) {
 		mgh := mocks.NewMockGitHub()
 
-		githubOrganization, githubRepository, _, githubUser, _ := mgh.Seed()
+		githubOrganization, githubRepository, _, githubUser, _, _ := mgh.Seed()
 
 		githubClient := github.NewClient(mgh.Server())
 		cache := newOrgNameCache(githubClient)

pkg/connector/team_test.go

@@ -20,7 +20,7 @@ func TestTeam(t *testing.T) {
 	t.Run("should grant and revoke entitlements", func(t *testing.T) {
 		mgh := mocks.NewMockGitHub()
 
-		githubOrganization, _, githubTeam, githubUser, _ := mgh.Seed()
+		githubOrganization, _, githubTeam, githubUser, _, _ := mgh.Seed()
 
 		githubClient := github.NewClient(mgh.Server())
 		cache := newOrgNameCache(githubClient)

pkg/connector/user_test.go

@@ -29,7 +29,7 @@ func TestUsersList(t *testing.T) {
 		t.Run(fmt.Sprintf("should get a list of users (SAML:%s)", testCase.message), func(t *testing.T) {
 			mgh := mocks.NewMockGitHub()
 
-			githubOrganization, _, _, githubUser, _ := mgh.Seed()
+			githubOrganization, _, _, githubUser, _, _ := mgh.Seed()
 
 			organization, err := organizationResource(
 				ctx,

test/mocks/github.go

@@ -93,6 +93,7 @@ func (mgh MockGitHub) Seed() (
 	*github.Repository,
 	*github.Team,
 	*github.User,
+	*OrganizationRole,
 	error,
 ) {
 	organizationId := int64(12)
@@ -137,7 +138,16 @@ func (mgh MockGitHub) Seed() (
 	mgh.teamMemberships[teamId] = mapset.NewSet[int64](userId)
 	mgh.organizationMemberships[organizationId] = mapset.NewSet[int64](userId)
 
-	return &githubOrganization, &githubRepository, &githubTeam, &githubUser, nil
+	// Add a mock org role
+	roleId := int64(1)
+	orgRole := &OrganizationRole{
+		ID:          roleId,
+		Name:        "Test Role",
+		Description: "Test Role Description",
+	}
+	mgh.orgRoles[roleId] = mapset.NewSet[int64]() // Initialize the set for this role
+
+	return &githubOrganization, &githubRepository, &githubTeam, &githubUser, orgRole, nil
 }
 
 func getResource[T interface{}](
@@ -564,13 +574,29 @@ func (mgh MockGitHub) getOrgRoleTeams(
 	start := (page - 1) * perPage
 	end := start + perPage
 	i := 0
-	for _, team := range mgh.teams {
-		if i >= start && i < end {
-			teams = append(teams, &team)
+
+	// Get users with this role
+	roleUsers := mgh.orgRoles[roleID]
+
+	// Find teams that have members with this role
+	for teamID, team := range mgh.teams {
+		teamMembers := mgh.teamMemberships[teamID]
+		if teamMembers == nil {
+			continue
 		}
-		i++
-		if i >= end {
-			break
+
+		// Check if any team member has the role
+		for _, userID := range teamMembers.ToSlice() {
+			if roleUsers.Contains(userID) {
+				if i >= start && i < end {
+					teams = append(teams, &team)
+				}
+				i++
+				if i >= end {
+					break
+				}
+				break // Found a member with the role, no need to check other members
+			}
 		}
 	}
 
@@ -647,6 +673,30 @@ func (mgh MockGitHub) removeOrgRoleUser(
 	}
 }
 
+// Add a handler for GET /orgs/{org}/organization-roles/{role_id}
+func (mgh MockGitHub) getOrgRoleByID(
+	w http.ResponseWriter,
+	variables map[string]string,
+) {
+	orgID, _ := getCrossTableId(w, variables, "org")
+	roleID, _ := getCrossTableId(w, variables, "role_id")
+	if _, ok := mgh.organizations[orgID]; !ok {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	// Only support role ID 1 for the mock
+	if roleID != 1 {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	role := &OrganizationRole{
+		ID:          1,
+		Name:        "Test Role",
+		Description: "Test Role Description",
+	}
+	_, _ = w.Write(mock.MustMarshal(role))
+}
+
 type handler = func(w http.ResponseWriter, variables map[string]string)
 
 // addEndpointHandler takes a string interpolation pattern and a handler
@@ -694,9 +744,13 @@ func (mgh MockGitHub) Server() *http.Client {
 		DeleteOrganizationsTeamsMembershipsByOrganizationByTeamIdByUsername: mgh.removeMembership,
 		PutOrganizationsTeamsMembershipsByOrganizationByTeamIdByUsername:    mgh.addMembership,
 		// Add organization role endpoints
-		GetOrgsRolesByOrg:                           mgh.getOrgRoles,
-		GetOrgsRolesTeamsByOrgByRoleId:              mgh.getOrgRoleTeams,
-		GetOrgsRolesUsersByOrgByRoleId:              mgh.getOrgRoleUsers,
+		GetOrgsRolesByOrg:              mgh.getOrgRoles,
+		GetOrgsRolesTeamsByOrgByRoleId: mgh.getOrgRoleTeams,
+		GetOrgsRolesUsersByOrgByRoleId: mgh.getOrgRoleUsers,
+		mock.EndpointPattern{
+			Pattern: "/orgs/{org}/organization-roles/{role_id}",
+			Method:  "GET",
+		}: mgh.getOrgRoleByID,
 		PutOrgsRolesUsersByOrgByRoleIdByUsername:    mgh.addOrgRoleUser,
 		DeleteOrgsRolesUsersByOrgByRoleIdByUsername: mgh.removeOrgRoleUser,
 	}
@@ -721,10 +775,10 @@ func (mgh *MockGitHub) AddUserToOrgRole(roleID int64, userID 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]()
+// AddMembership adds a user to a team for testing purposes.
+func (mgh *MockGitHub) AddMembership(teamID int64, userID int64) {
+	if _, ok := mgh.teamMemberships[teamID]; !ok {
+		mgh.teamMemberships[teamID] = mapset.NewSet[int64]()
 	}
-	mgh.orgRoles[roleID].Add(teamID)
+	mgh.teamMemberships[teamID].Add(userID)
 }