-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Expand file tree
/
Copy pathai_api_test.go
More file actions
250 lines (218 loc) · 6.91 KB
/
ai_api_test.go
File metadata and controls
250 lines (218 loc) · 6.91 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
// Copyright 2026 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package casbin
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
// TestExplainWithoutConfig tests that Explain returns error when config is not set.
func TestExplainWithoutConfig(t *testing.T) {
e, err := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")
if err != nil {
t.Fatal(err)
}
_, err = e.Explain("alice", "data1", "read")
if err == nil {
t.Error("Expected error when AI config is not set")
}
if !strings.Contains(err.Error(), "AI config not set") {
t.Errorf("Expected 'AI config not set' error, got: %v", err)
}
}
// TestExplainWithMockAPI tests Explain with a mock OpenAI-compatible API.
func TestExplainWithMockAPI(t *testing.T) {
// Create a mock server that simulates OpenAI API
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Verify request
if r.Method != http.MethodPost {
t.Errorf("Expected POST request, got %s", r.Method)
}
if r.Header.Get("Content-Type") != "application/json" {
t.Errorf("Expected Content-Type: application/json, got %s", r.Header.Get("Content-Type"))
}
if !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") {
t.Errorf("Expected Bearer token in Authorization header, got %s", r.Header.Get("Authorization"))
}
// Parse request to verify structure
var req aiChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Errorf("Failed to decode request: %v", err)
}
if req.Model != "gpt-3.5-turbo" {
t.Errorf("Expected model gpt-3.5-turbo, got %s", req.Model)
}
if len(req.Messages) != 2 {
t.Errorf("Expected 2 messages, got %d", len(req.Messages))
}
// Send mock response
resp := aiChatResponse{
Choices: []struct {
Message aiMessage `json:"message"`
}{
{
Message: aiMessage{
Role: "assistant",
Content: "The request was allowed because alice has read permission on data1 according to the policy rule.",
},
},
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(resp)
}))
defer mockServer.Close()
// Create enforcer
e, err := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")
if err != nil {
t.Fatal(err)
}
// Set AI config with mock server
e.SetAIConfig(AIConfig{
Endpoint: mockServer.URL,
APIKey: "test-api-key",
Model: "gpt-3.5-turbo",
Timeout: 5 * time.Second,
})
// Test explanation for allowed request
explanation, err := e.Explain("alice", "data1", "read")
if err != nil {
t.Fatalf("Failed to get explanation: %v", err)
}
if explanation == "" {
t.Error("Expected non-empty explanation")
}
if !strings.Contains(explanation, "allowed") {
t.Errorf("Expected explanation to mention 'allowed', got: %s", explanation)
}
}
// TestExplainDenied tests Explain for a denied request.
func TestExplainDenied(t *testing.T) {
// Create a mock server
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
resp := aiChatResponse{
Choices: []struct {
Message aiMessage `json:"message"`
}{
{
Message: aiMessage{
Role: "assistant",
Content: "The request was denied because there is no policy rule that allows alice to write to data1.",
},
},
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(resp)
}))
defer mockServer.Close()
// Create enforcer
e, err := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")
if err != nil {
t.Fatal(err)
}
// Set AI config
e.SetAIConfig(AIConfig{
Endpoint: mockServer.URL,
APIKey: "test-api-key",
Model: "gpt-3.5-turbo",
Timeout: 5 * time.Second,
})
// Test explanation for denied request
explanation, err := e.Explain("alice", "data1", "write")
if err != nil {
t.Fatalf("Failed to get explanation: %v", err)
}
if explanation == "" {
t.Error("Expected non-empty explanation")
}
if !strings.Contains(explanation, "denied") {
t.Errorf("Expected explanation to mention 'denied', got: %s", explanation)
}
}
// TestExplainAPIError tests handling of API errors.
func TestExplainAPIError(t *testing.T) {
// Create a mock server that returns an error
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
resp := aiChatResponse{
Error: &struct {
Message string `json:"message"`
}{
Message: "Invalid API key",
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusUnauthorized)
json.NewEncoder(w).Encode(resp)
}))
defer mockServer.Close()
// Create enforcer
e, err := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")
if err != nil {
t.Fatal(err)
}
// Set AI config
e.SetAIConfig(AIConfig{
Endpoint: mockServer.URL,
APIKey: "invalid-key",
Model: "gpt-3.5-turbo",
Timeout: 5 * time.Second,
})
// Test that API error is properly handled
_, err = e.Explain("alice", "data1", "read")
if err == nil {
t.Error("Expected error for API failure")
}
if !strings.Contains(err.Error(), "Invalid API key") {
t.Errorf("Expected API error message, got: %v", err)
}
}
// TestBuildExplainContext tests the context building function.
func TestBuildExplainContext(t *testing.T) {
e, err := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")
if err != nil {
t.Fatal(err)
}
// Test with matched rules
rvals := []interface{}{"alice", "data1", "read"}
result := true
matchedRules := []string{"alice, data1, read"}
context := e.buildExplainContext(rvals, result, matchedRules)
// Verify context contains expected elements
if !strings.Contains(context, "alice") {
t.Error("Context should contain subject 'alice'")
}
if !strings.Contains(context, "data1") {
t.Error("Context should contain object 'data1'")
}
if !strings.Contains(context, "read") {
t.Error("Context should contain action 'read'")
}
if !strings.Contains(context, "true") {
t.Error("Context should contain result 'true'")
}
if !strings.Contains(context, "alice, data1, read") {
t.Error("Context should contain matched rule")
}
// Test with no matched rules
context2 := e.buildExplainContext(rvals, false, []string{})
if !strings.Contains(context2, "No policy rules matched") {
t.Error("Context should indicate no matched rules")
}
}