From 2081248b1b55baee8221289a2c401c96c0e85aea Mon Sep 17 00:00:00 2001
From: Stamatis Katsaounis <katsaouniss@gmail.com>
Date: Fri, 12 May 2023 11:12:33 +0300
Subject: [PATCH] Add GetRoles call to fetch all Grafana roles

---
 role.go      | 27 ++++++++++++++++
 role_test.go | 90 ++++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 110 insertions(+), 7 deletions(-)

diff --git a/role.go b/role.go
index 53236007..5fec60c1 100644
--- a/role.go
+++ b/role.go
@@ -3,6 +3,7 @@ package gapi
 import (
 	"encoding/json"
 	"fmt"
+	"net/url"
 )
 
 type Role struct {
@@ -22,6 +23,32 @@ type Permission struct {
 	Scope  string `json:"scope"`
 }
 
+// GetRole fetches and returns Grafana roles. Available only in Grafana Enterprise 8.+.
+func (c *Client) GetRoles() ([]Role, error) {
+	const limit = 1000
+	var (
+		page     = 0
+		newRoles []Role
+		roles    []Role
+		query    = make(url.Values)
+	)
+	query.Set("limit", fmt.Sprint(limit))
+	for {
+		page++
+		query.Set("page", fmt.Sprint(page))
+
+		if err := c.request("GET", "/api/access-control/roles", query, nil, &newRoles); err != nil {
+			return nil, err
+		}
+
+		roles = append(roles, newRoles...)
+
+		if len(newRoles) < limit {
+			return roles, nil
+		}
+	}
+}
+
 // GetRole gets a role with permissions for the given UID. Available only in Grafana Enterprise 8.+.
 func (c *Client) GetRole(uid string) (*Role, error) {
 	r := &Role{}
diff --git a/role_test.go b/role_test.go
index 3e0debcd..f5c1fae7 100644
--- a/role_test.go
+++ b/role_test.go
@@ -1,6 +1,7 @@
 package gapi
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/gobs/pretty"
@@ -55,8 +56,86 @@ const (
 
 	updatedRoleResponse = `{"message":"Role updated"}`
 	deleteRoleResponse  = `{"message":"Role deleted"}`
+
+	roleUID = "vc3SCSsGz"
 )
 
+func TestRoles(t *testing.T) {
+	mockData := strings.Repeat(getRoleResponse+",", 1000) // make 1000 roles.
+	mockData = "[" + mockData[:len(mockData)-1] + "]"     // remove trailing comma; make a json list.
+
+	// This creates 1000 + 1000 + 1 (2001, 3 calls) worth of roles.
+	client := gapiTestToolsFromCalls(t, []mockServerCall{
+		{200, mockData},
+		{200, mockData},
+		{200, "[" + getRoleResponse + "]"},
+	})
+
+	const dashCount = 2001
+
+	roles, err := client.GetRoles()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log(pretty.PrettyFormat(roles))
+
+	if len(roles) != dashCount {
+		t.Fatalf("Length of returned roles should be %d", dashCount)
+	}
+	if roles[0].UID != roleUID || roles[0].Name != "test:policy" {
+		t.Error("Not correctly parsing returned roles.")
+	}
+	if roles[dashCount-1].UID != roleUID || roles[dashCount-1].Name != "test:policy" {
+		t.Error("Not correctly parsing returned roles.")
+	}
+}
+
+func TestRolesZeroResults(t *testing.T) {
+	// This return zero roles.
+	client := gapiTestToolsFromCalls(t, []mockServerCall{
+		{200, "[]"},
+	})
+
+	roles, err := client.GetRoles()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(roles) != 0 {
+		t.Errorf("Length of returned roles should be zero")
+	}
+}
+
+func TestRolesSinglePage(t *testing.T) {
+	mockData := strings.Repeat(getRoleResponse+",", 999) // make 999 roles.
+	mockData = "[" + mockData[:len(mockData)-1] + "]"    // remove trailing comma; make a json list.
+
+	// This creates 999 worth of roles.
+	client := gapiTestToolsFromCalls(t, []mockServerCall{
+		{200, mockData},
+	})
+
+	const dashCount = 999
+
+	roles, err := client.GetRoles()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log(pretty.PrettyFormat(roles))
+
+	if len(roles) != dashCount {
+		t.Fatalf("Length of returned roles should be %d", dashCount)
+	}
+	if roles[0].UID != roleUID || roles[0].Name != "test:policy" {
+		t.Error("Not correctly parsing returned roles.")
+	}
+	if roles[dashCount-1].UID != roleUID || roles[dashCount-1].Name != "test:policy" {
+		t.Error("Not correctly parsing returned roles.")
+	}
+}
+
 func TestNewRole(t *testing.T) {
 	client := gapiTestTools(t, 201, newRoleResponse)
 
@@ -79,7 +158,7 @@ func TestNewRole(t *testing.T) {
 
 	t.Log(pretty.PrettyFormat(resp))
 
-	if resp.UID != "vc3SCSsGz" {
+	if resp.UID != roleUID {
 		t.Error("Not correctly parsing returned role uid.")
 	}
 }
@@ -87,10 +166,7 @@ func TestNewRole(t *testing.T) {
 func TestGetRole(t *testing.T) {
 	client := gapiTestTools(t, 200, getRoleResponse)
 
-	uid := "vc3SCSsGz"
-
-	resp, err := client.GetRole(uid)
-
+	resp, err := client.GetRole(roleUID)
 	if err != nil {
 		t.Error(err)
 	}
@@ -98,7 +174,7 @@ func TestGetRole(t *testing.T) {
 	expected := Role{
 		Global:      false,
 		Version:     1,
-		UID:         "vc3SCSsGz",
+		UID:         roleUID,
 		Name:        "test:policy",
 		Description: "Test policy description",
 		Group:       "test group",
@@ -143,7 +219,7 @@ func TestUpdateRole(t *testing.T) {
 func TestDeleteRole(t *testing.T) {
 	client := gapiTestTools(t, 200, deleteRoleResponse)
 
-	err := client.DeleteRole("vc3SCSsGz", false)
+	err := client.DeleteRole(roleUID, false)
 	if err != nil {
 		t.Error(err)
 	}