Extend user.SessionState CustomPolicies API, for tests

user/custom_policies.go

@@ -6,10 +6,26 @@ import (
 	"fmt"
 )
 
+// CustomPolicies returns a map of custom policies on the session.
+// To preserve policy order, use GetCustomPolicies instead.
 func (s *SessionState) CustomPolicies() (map[string]Policy, error) {
+	customPolicies, err := s.GetCustomPolicies()
+	if err != nil {
+		return nil, err
+	}
+
+	result := make(map[string]Policy, len(customPolicies))
+	for i := 0; i < len(customPolicies); i++ {
+		result[customPolicies[i].ID] = customPolicies[i]
+	}
+
+	return result, nil
+}
+
+// GetCustomPolicies is like CustomPolicies but returns the list, preserving order.
+func (s *SessionState) GetCustomPolicies() ([]Policy, error) {
 	var (
 		customPolicies []Policy
-		ret            map[string]Policy
 	)
 
 	metadataPolicies, found := s.MetaData["policies"].([]interface{})
@@ -22,16 +38,14 @@ func (s *SessionState) CustomPolicies() (map[string]Policy, error) {
 		return nil, fmt.Errorf("failed to marshal metadata policies: %w", err)
 	}
 
-	_ = json.Unmarshal(polJSON, &customPolicies)
-
-	ret = make(map[string]Policy, len(customPolicies))
-	for i := 0; i < len(customPolicies); i++ {
-		ret[customPolicies[i].ID] = customPolicies[i]
+	if err := json.Unmarshal(polJSON, &customPolicies); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal metadata policies: %w", err)
 	}
 
-	return ret, nil
+	return customPolicies, err
 }
 
+// SetCustomPolicies sets custom policies into session metadata.
 func (s *SessionState) SetCustomPolicies(list []Policy) {
 	if s.MetaData == nil {
 		s.MetaData = make(map[string]interface{})