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