forge_core/testing/context/
query.rs

1//! Test context for query functions.
2
3#![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
16/// Test context for query functions with optional DB access.
17pub struct TestQueryContext {
18    pub auth: AuthContext,
19    pub request: RequestMetadata,
20    pool: Option<PgPool>,
21    tenant_id: Option<Uuid>,
22    env_provider: Arc<MockEnvProvider>,
23}
24
25impl TestQueryContext {
26    /// Create a new builder.
27    pub fn builder() -> TestQueryContextBuilder {
28        TestQueryContextBuilder::default()
29    }
30
31    /// Create a minimal unauthenticated context (no database).
32    pub fn minimal() -> Self {
33        Self::builder().build()
34    }
35
36    /// Create an authenticated context with the given user ID (no database).
37    pub fn authenticated(user_id: Uuid) -> Self {
38        Self::builder().as_user(user_id).build()
39    }
40
41    /// Create a context with a database pool.
42    pub fn with_pool(pool: PgPool, user_id: Option<Uuid>) -> Self {
43        let mut builder = Self::builder().with_pool(pool);
44        if let Some(id) = user_id {
45            builder = builder.as_user(id);
46        }
47        builder.build()
48    }
49
50    /// Get the database pool (if available).
51    pub fn db(&self) -> Option<&PgPool> {
52        self.pool.as_ref()
53    }
54
55    /// Get the authenticated user's UUID. Returns 401 if not authenticated.
56    pub fn user_id(&self) -> Result<Uuid> {
57        self.auth.require_user_id()
58    }
59
60    /// Check if a specific role is present.
61    pub fn has_role(&self, role: &str) -> bool {
62        self.auth.has_role(role)
63    }
64
65    /// Get a claim value.
66    pub fn claim(&self, key: &str) -> Option<&serde_json::Value> {
67        self.auth.claim(key)
68    }
69
70    /// Get the tenant ID (if set).
71    pub fn tenant_id(&self) -> Option<Uuid> {
72        self.tenant_id
73    }
74
75    /// Get the mock env provider for verification.
76    pub fn env_mock(&self) -> &MockEnvProvider {
77        &self.env_provider
78    }
79}
80
81impl EnvAccess for TestQueryContext {
82    fn env_provider(&self) -> &dyn EnvProvider {
83        self.env_provider.as_ref()
84    }
85}
86
87/// Builder for TestQueryContext.
88#[derive(Default)]
89pub struct TestQueryContextBuilder {
90    user_id: Option<Uuid>,
91    roles: Vec<String>,
92    claims: HashMap<String, serde_json::Value>,
93    tenant_id: Option<Uuid>,
94    pool: Option<PgPool>,
95    env_vars: HashMap<String, String>,
96}
97
98impl TestQueryContextBuilder {
99    /// Set the authenticated user with a UUID.
100    pub fn as_user(mut self, id: Uuid) -> Self {
101        self.user_id = Some(id);
102        self
103    }
104
105    /// For non-UUID auth providers (Firebase, Clerk, etc.).
106    pub fn as_subject(mut self, subject: impl Into<String>) -> Self {
107        self.claims
108            .insert("sub".to_string(), serde_json::json!(subject.into()));
109        self
110    }
111
112    /// Add a role.
113    pub fn with_role(mut self, role: impl Into<String>) -> Self {
114        self.roles.push(role.into());
115        self
116    }
117
118    /// Add multiple roles.
119    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
120        self.roles.extend(roles);
121        self
122    }
123
124    /// Add a custom claim.
125    pub fn with_claim(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
126        self.claims.insert(key.into(), value);
127        self
128    }
129
130    /// Set the tenant ID for multi-tenant testing.
131    pub fn with_tenant(mut self, tenant_id: Uuid) -> Self {
132        self.tenant_id = Some(tenant_id);
133        self
134    }
135
136    /// Set the database pool.
137    pub fn with_pool(mut self, pool: PgPool) -> Self {
138        self.pool = Some(pool);
139        self
140    }
141
142    /// Set a single environment variable.
143    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
144        self.env_vars.insert(key.into(), value.into());
145        self
146    }
147
148    /// Set multiple environment variables.
149    pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
150        self.env_vars.extend(vars);
151        self
152    }
153
154    /// Build the test context.
155    pub fn build(self) -> TestQueryContext {
156        TestQueryContext {
157            auth: build_test_auth(self.user_id, self.roles, self.claims),
158            request: RequestMetadata::default(),
159            pool: self.pool,
160            tenant_id: self.tenant_id,
161            env_provider: Arc::new(MockEnvProvider::with_vars(self.env_vars)),
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn test_minimal_context() {
172        let ctx = TestQueryContext::minimal();
173        assert!(!ctx.auth.is_authenticated());
174        assert!(ctx.db().is_none());
175    }
176
177    #[test]
178    fn test_authenticated_context() {
179        let user_id = Uuid::new_v4();
180        let ctx = TestQueryContext::authenticated(user_id);
181        assert!(ctx.auth.is_authenticated());
182        assert_eq!(ctx.user_id().unwrap(), user_id);
183    }
184
185    #[test]
186    fn test_context_with_roles() {
187        let ctx = TestQueryContext::builder()
188            .as_user(Uuid::new_v4())
189            .with_role("admin")
190            .with_role("user")
191            .build();
192
193        assert!(ctx.has_role("admin"));
194        assert!(ctx.has_role("user"));
195        assert!(!ctx.has_role("superuser"));
196    }
197
198    #[test]
199    fn test_context_with_claims() {
200        let ctx = TestQueryContext::builder()
201            .as_user(Uuid::new_v4())
202            .with_claim("org_id", serde_json::json!("org-123"))
203            .build();
204
205        assert_eq!(ctx.claim("org_id"), Some(&serde_json::json!("org-123")));
206        assert!(ctx.claim("nonexistent").is_none());
207    }
208
209    #[test]
210    fn test_context_with_env() {
211        let ctx = TestQueryContext::builder()
212            .with_env("API_KEY", "test_key_123")
213            .with_env("TIMEOUT", "30")
214            .build();
215
216        // Test env access via EnvAccess trait
217        assert_eq!(ctx.env("API_KEY"), Some("test_key_123".to_string()));
218        assert_eq!(ctx.env_or("TIMEOUT", "10"), "30");
219        assert_eq!(ctx.env_or("MISSING", "default"), "default");
220
221        // Test env_require
222        assert!(ctx.env_require("API_KEY").is_ok());
223        assert!(ctx.env_require("MISSING").is_err());
224
225        // Test env_parse
226        let timeout: u32 = ctx.env_parse("TIMEOUT").unwrap();
227        assert_eq!(timeout, 30);
228
229        // Verify access tracking
230        ctx.env_mock().assert_accessed("API_KEY");
231        ctx.env_mock().assert_accessed("TIMEOUT");
232    }
233}
forge_core/testing/context/query.rs

forge_core/testing/context/
query.rs
