Skip to main content

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.
17///
18/// Provides an isolated testing environment for queries with configurable
19/// authentication and optional database access.
20///
21/// # Example
22///
23/// ```ignore
24/// // Simple authenticated context
25/// let ctx = TestQueryContext::authenticated(Uuid::new_v4());
26///
27/// // Context with roles and claims
28/// let ctx = TestQueryContext::builder()
29///     .as_user(Uuid::new_v4())
30///     .with_role("admin")
31///     .with_claim("org_id", json!("org-123"))
32///     .build();
33/// ```
34pub struct TestQueryContext {
35    /// Authentication context.
36    pub auth: AuthContext,
37    /// Request metadata.
38    pub request: RequestMetadata,
39    /// Optional database pool for integration tests.
40    pool: Option<PgPool>,
41    /// Tenant ID for multi-tenant testing.
42    tenant_id: Option<Uuid>,
43    /// Mock environment provider.
44    env_provider: Arc<MockEnvProvider>,
45}
46
47impl TestQueryContext {
48    /// Create a new builder.
49    pub fn builder() -> TestQueryContextBuilder {
50        TestQueryContextBuilder::default()
51    }
52
53    /// Create a minimal unauthenticated context (no database).
54    pub fn minimal() -> Self {
55        Self::builder().build()
56    }
57
58    /// Create an authenticated context with the given user ID (no database).
59    pub fn authenticated(user_id: Uuid) -> Self {
60        Self::builder().as_user(user_id).build()
61    }
62
63    /// Create a context with a database pool.
64    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    /// Get the database pool (if available).
73    pub fn db(&self) -> Option<&PgPool> {
74        self.pool.as_ref()
75    }
76
77    /// Get the authenticated user ID or return an error.
78    pub fn require_user_id(&self) -> Result<Uuid> {
79        self.auth.require_user_id()
80    }
81
82    /// Like `require_user_id()` but for non-UUID auth providers.
83    pub fn require_subject(&self) -> Result<&str> {
84        self.auth.require_subject()
85    }
86
87    /// Check if a specific role is present.
88    pub fn has_role(&self, role: &str) -> bool {
89        self.auth.has_role(role)
90    }
91
92    /// Get a claim value.
93    pub fn claim(&self, key: &str) -> Option<&serde_json::Value> {
94        self.auth.claim(key)
95    }
96
97    /// Get the tenant ID (if set).
98    pub fn tenant_id(&self) -> Option<Uuid> {
99        self.tenant_id
100    }
101
102    /// Get the mock env provider for verification.
103    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/// Builder for TestQueryContext.
115#[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    /// Set the authenticated user with a UUID.
127    pub fn as_user(mut self, id: Uuid) -> Self {
128        self.user_id = Some(id);
129        self
130    }
131
132    /// For non-UUID auth providers (Firebase, Clerk, etc.).
133    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    /// Add a role.
140    pub fn with_role(mut self, role: impl Into<String>) -> Self {
141        self.roles.push(role.into());
142        self
143    }
144
145    /// Add multiple roles.
146    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
147        self.roles.extend(roles);
148        self
149    }
150
151    /// Add a custom claim.
152    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    /// Set the tenant ID for multi-tenant testing.
158    pub fn with_tenant(mut self, tenant_id: Uuid) -> Self {
159        self.tenant_id = Some(tenant_id);
160        self
161    }
162
163    /// Set the database pool.
164    pub fn with_pool(mut self, pool: PgPool) -> Self {
165        self.pool = Some(pool);
166        self
167    }
168
169    /// Set a single environment variable.
170    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    /// Set multiple environment variables.
176    pub fn with_envs(mut self, vars: HashMap<String, String>) -> Self {
177        self.env_vars.extend(vars);
178        self
179    }
180
181    /// Build the test context.
182    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        // Test env access via EnvAccess trait
244        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        // Test env_require
249        assert!(ctx.env_require("API_KEY").is_ok());
250        assert!(ctx.env_require("MISSING").is_err());
251
252        // Test env_parse
253        let timeout: u32 = ctx.env_parse("TIMEOUT").unwrap();
254        assert_eq!(timeout, 30);
255
256        // Verify access tracking
257        ctx.env_mock().assert_accessed("API_KEY");
258        ctx.env_mock().assert_accessed("TIMEOUT");
259    }
260}