Skip to main content

forge_core/testing/context/
query.rs

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