forge_core/testing/context/
query.rs1use std::collections::HashMap;
4use std::sync::Arc;
5
6use sqlx::PgPool;
7use uuid::Uuid;
8
9use super::build_test_auth;
10use crate::Result;
11use crate::env::{EnvAccess, EnvProvider, MockEnvProvider};
12use crate::function::{AuthContext, RequestMetadata};
13
14pub struct TestQueryContext {
33 pub auth: AuthContext,
35 pub request: RequestMetadata,
37 pool: Option<PgPool>,
39 tenant_id: Option<Uuid>,
41 env_provider: Arc<MockEnvProvider>,
43}
44
45impl TestQueryContext {
46 pub fn builder() -> TestQueryContextBuilder {
48 TestQueryContextBuilder::default()
49 }
50
51 pub fn minimal() -> Self {
53 Self::builder().build()
54 }
55
56 pub fn authenticated(user_id: Uuid) -> Self {
58 Self::builder().as_user(user_id).build()
59 }
60
61 pub fn with_pool(pool: PgPool, user_id: Option<Uuid>) -> Self {
63 let mut builder = Self::builder().with_pool(pool);
64 if let Some(id) = user_id {
65 builder = builder.as_user(id);
66 }
67 builder.build()
68 }
69
70 pub fn db(&self) -> Option<&PgPool> {
72 self.pool.as_ref()
73 }
74
75 pub fn require_user_id(&self) -> Result<Uuid> {
77 self.auth.require_user_id()
78 }
79
80 pub fn require_subject(&self) -> Result<&str> {
82 self.auth.require_subject()
83 }
84
85 pub fn has_role(&self, role: &str) -> bool {
87 self.auth.has_role(role)
88 }
89
90 pub fn claim(&self, key: &str) -> Option<&serde_json::Value> {
92 self.auth.claim(key)
93 }
94
95 pub fn tenant_id(&self) -> Option<Uuid> {
97 self.tenant_id
98 }
99
100 pub fn env_mock(&self) -> &MockEnvProvider {
102 &self.env_provider
103 }
104}
105
106impl EnvAccess for TestQueryContext {
107 fn env_provider(&self) -> &dyn EnvProvider {
108 self.env_provider.as_ref()
109 }
110}
111
112#[derive(Default)]
114pub struct TestQueryContextBuilder {
115 user_id: Option<Uuid>,
116 roles: Vec<String>,
117 claims: HashMap<String, serde_json::Value>,
118 tenant_id: Option<Uuid>,
119 pool: Option<PgPool>,
120 env_vars: HashMap<String, String>,
121}
122
123impl TestQueryContextBuilder {
124 pub fn as_user(mut self, id: Uuid) -> Self {
126 self.user_id = Some(id);
127 self
128 }
129
130 pub fn as_subject(mut self, subject: impl Into<String>) -> Self {
132 self.claims
133 .insert("sub".to_string(), serde_json::json!(subject.into()));
134 self
135 }
136
137 pub fn with_role(mut self, role: impl Into<String>) -> Self {
139 self.roles.push(role.into());
140 self
141 }
142
143 pub fn with_roles(mut self, roles: Vec<String>) -> Self {
145 self.roles.extend(roles);
146 self
147 }
148
149 pub fn with_claim(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
151 self.claims.insert(key.into(), value);
152 self
153 }
154
155 pub fn with_tenant(mut self, tenant_id: Uuid) -> Self {
157 self.tenant_id = Some(tenant_id);
158 self
159 }
160
161 pub fn with_pool(mut self, pool: PgPool) -> Self {
163 self.pool = Some(pool);
164 self
165 }
166
167 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
169 self.env_vars.insert(key.into(), value.into());
170 self
171 }
172
173 pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
175 self.env_vars.extend(vars);
176 self
177 }
178
179 pub fn build(self) -> TestQueryContext {
181 TestQueryContext {
182 auth: build_test_auth(self.user_id, self.roles, self.claims),
183 request: RequestMetadata::default(),
184 pool: self.pool,
185 tenant_id: self.tenant_id,
186 env_provider: Arc::new(MockEnvProvider::with_vars(self.env_vars)),
187 }
188 }
189}
190
191#[cfg(test)]
192mod tests {
193 use super::*;
194
195 #[test]
196 fn test_minimal_context() {
197 let ctx = TestQueryContext::minimal();
198 assert!(!ctx.auth.is_authenticated());
199 assert!(ctx.db().is_none());
200 }
201
202 #[test]
203 fn test_authenticated_context() {
204 let user_id = Uuid::new_v4();
205 let ctx = TestQueryContext::authenticated(user_id);
206 assert!(ctx.auth.is_authenticated());
207 assert_eq!(ctx.require_user_id().unwrap(), user_id);
208 }
209
210 #[test]
211 fn test_context_with_roles() {
212 let ctx = TestQueryContext::builder()
213 .as_user(Uuid::new_v4())
214 .with_role("admin")
215 .with_role("user")
216 .build();
217
218 assert!(ctx.has_role("admin"));
219 assert!(ctx.has_role("user"));
220 assert!(!ctx.has_role("superuser"));
221 }
222
223 #[test]
224 fn test_context_with_claims() {
225 let ctx = TestQueryContext::builder()
226 .as_user(Uuid::new_v4())
227 .with_claim("org_id", serde_json::json!("org-123"))
228 .build();
229
230 assert_eq!(ctx.claim("org_id"), Some(&serde_json::json!("org-123")));
231 assert!(ctx.claim("nonexistent").is_none());
232 }
233
234 #[test]
235 fn test_context_with_env() {
236 let ctx = TestQueryContext::builder()
237 .with_env("API_KEY", "test_key_123")
238 .with_env("TIMEOUT", "30")
239 .build();
240
241 assert_eq!(ctx.env("API_KEY"), Some("test_key_123".to_string()));
243 assert_eq!(ctx.env_or("TIMEOUT", "10"), "30");
244 assert_eq!(ctx.env_or("MISSING", "default"), "default");
245
246 assert!(ctx.env_require("API_KEY").is_ok());
248 assert!(ctx.env_require("MISSING").is_err());
249
250 let timeout: u32 = ctx.env_parse("TIMEOUT").unwrap();
252 assert_eq!(timeout, 30);
253
254 ctx.env_mock().assert_accessed("API_KEY");
256 ctx.env_mock().assert_accessed("TIMEOUT");
257 }
258}