added test for GetAllowedUsersViaPgxConnection

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
Dmitri Dolguikh
2026-08-05 16:34:44 +02:00
parent ba574dc739
commit 07809fe923
2 changed files with 66 additions and 8 deletions

View File

@@ -0,0 +1,62 @@
package networkmap_pgsql
import (
"context"
"testing"
networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
"github.com/stretchr/testify/assert"
)
func TestGetAllowedUsers(t *testing.T) {
ctx := context.TODO()
_, err := pgstore.Pool.Query(ctx,
`insert into users (id, name, account_id, auto_groups, blocked, is_service_user)
VALUES('user-1','user-1','account-1','["group-one-resource-id"]',false,false)`)
assert.NoError(t, err)
_, err = pgstore.Pool.Query(ctx,
`insert into users (id, name, account_id, auto_groups, blocked, is_service_user)
VALUES('user-2','user-2','account-1','["group-one-resource-id","group-two-resources-id"]',false,false)`)
assert.NoError(t, err)
_, err = pgstore.Pool.Query(ctx,
`insert into users (id, name, account_id, auto_groups, blocked, is_service_user)
VALUES('user-3','user-3','account-1','["group-two-resources-id"]',false,false)`)
// shouldn't be included as it's blocked
_, err = pgstore.Pool.Query(ctx,
`insert into users (id, name, account_id, auto_groups, blocked, is_service_user)
VALUES('user-3','user-3','account-1','["group-two-resources-id"]',true,false)`)
assert.NoError(t, err)
// shouldn't be included as it's a service_user
_, err = pgstore.Pool.Query(ctx,
`insert into users (id, name, account_id, auto_groups, blocked, is_service_user)
VALUES('user-3','user-3','account-1','["group-two-resources-id"]',false,true)`)
_, err = pgstore.Pool.Query(ctx,
`insert into groups (id, name, account_id)
VALUES('all-group-1','All','account-1')`)
assert.NoError(t, err)
_, err = pgstore.Pool.Query(ctx,
`insert into groups (id, name, account_id)
VALUES('all-group-2','All','account-1')`)
assert.NoError(t, err)
_, err = pgstore.Pool.Query(ctx,
`insert into groups (id, name, account_id)
VALUES('all-group-3','All','account-1')`)
assert.NoError(t, err)
userIdx, groupIdToUserIds, err := networkmap_pgsql.GetAllowedUsersViaPgxConnection(ctx, conn(t, ctx), "account-1")
assert.NoError(t, err)
assert.Equal(t, userIdx, map[string]struct{}{
"user-1": {},
"user-2": {},
"user-3": {},
})
assert.Equal(t, groupIdToUserIds, map[string][]string{
"group-one-resource-id": {"user-1", "user-2"},
"group-two-resources-id": {"user-2", "user-3"},
"all-group-1": {"user-1", "user-2", "user-3"},
"all-group-2": {"user-1", "user-2", "user-3"},
"all-group-3": {"user-1", "user-2", "user-3"},
})
}

View File

@@ -14,7 +14,7 @@ const (
`
GetAllGroupIdQuery = `
select id from groups
select array_agg(id) from groups
where account_id=$1 and name='All'
`
)
@@ -42,14 +42,10 @@ func GetAllowedUsersViaPgxConnection(ctx context.Context, con *pgx.Conn, account
if err != nil {
return nil, nil, err
}
allGroupIds, err := pgx.CollectRows(rows, pgx.RowTo[string])
allGroupIds, err := pgx.CollectOneRow(rows, pgx.RowTo[[]string])
if err != nil {
return nil, nil, err
}
allGroupId := ""
if len(allGroupIds) > 0 {
allGroupId = allGroupIds[0]
}
userIdIdx := make(map[string]struct{})
groupIdToUserIds := make(map[string][]string)
@@ -58,8 +54,8 @@ func GetAllowedUsersViaPgxConnection(ctx context.Context, con *pgx.Conn, account
for _, groupId := range user.AutoGroups {
groupIdToUserIds[groupId] = append(groupIdToUserIds[groupId], user.ID)
}
if allGroupId != "" {
groupIdToUserIds[allGroupId] = append(groupIdToUserIds[allGroupId], user.ID)
for _, allgid := range allGroupIds {
groupIdToUserIds[allgid] = append(groupIdToUserIds[allgid], user.ID)
}
}