Skip to main content

nntp_proxy/session/
core.rs

1//! Client session struct definition and builder pattern
2//!
3//! This module provides the core data structure for representing an active NNTP client session,
4//! along with a builder pattern for flexible session construction.
5
6use std::sync::Arc;
7
8use crate::auth::AuthHandler;
9use crate::config::RoutingMode;
10use crate::metrics::MetricsCollector;
11use crate::pool::BufferPool;
12use crate::router::BackendSelector;
13use crate::types::{ClientAddress, ClientId};
14
15use super::{AuthState, ModeState, SessionMode};
16
17/// Represents an active client session
18pub struct ClientSession {
19    pub(super) client_addr: ClientAddress,
20    pub(super) buffer_pool: BufferPool,
21    /// Unique identifier for this client
22    pub(super) client_id: ClientId,
23    /// Optional router for per-command routing mode
24    pub(super) router: Option<Arc<BackendSelector>>,
25    /// Session mode state (encapsulates mode and routing mode)
26    pub(super) mode_state: ModeState,
27    /// Authentication handler
28    pub(super) auth_handler: Arc<AuthHandler>,
29    /// Authentication state (encapsulates auth status and username)
30    pub(super) auth_state: AuthState,
31    /// Metrics collector for session statistics (always enabled)
32    pub(super) metrics: crate::metrics::MetricsCollector,
33
34    /// Connection statistics aggregator for logging connection creation
35    pub(super) connection_stats: Option<crate::metrics::ConnectionStatsAggregator>,
36
37    /// Article cache (always present - fixed-size, memory-backed, hybrid, or disabled)
38    pub(super) cache: Arc<crate::cache::UnifiedCache>,
39
40    /// Whether to cache article bodies (config-driven)
41    pub(super) cache_articles: bool,
42
43    /// Whether to use adaptive availability prechecking for STAT/HEAD
44    pub(super) adaptive_precheck: bool,
45}
46
47/// Builder for constructing `ClientSession` instances
48///
49/// Provides a fluent API for creating client sessions with different routing modes.
50///
51/// # Examples
52///
53/// ```
54/// use std::net::SocketAddr;
55/// use std::sync::Arc;
56/// use nntp_proxy::session::ClientSession;
57/// use nntp_proxy::pool::BufferPool;
58/// use nntp_proxy::router::BackendSelector;
59/// use nntp_proxy::config::RoutingMode;
60/// use nntp_proxy::types::BufferSize;
61/// use nntp_proxy::auth::AuthHandler;
62/// use nntp_proxy::metrics::MetricsCollector;
63///
64/// let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
65/// let buffer_pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 10);
66/// let auth_handler = Arc::new(AuthHandler::new(None, None).unwrap());
67/// let metrics = MetricsCollector::new(1);
68///
69/// // Stateful 1:1 routing mode
70/// let session = ClientSession::builder(addr.into(), buffer_pool.clone(), auth_handler.clone(), metrics.clone())
71///     .build();
72///
73/// // Per-command routing mode
74/// let router = Arc::new(BackendSelector::new());
75/// let session = ClientSession::builder(addr.into(), buffer_pool.clone(), auth_handler, metrics)
76///     .with_router(router)
77///     .with_routing_mode(RoutingMode::PerCommand)
78///     .build();
79/// ```
80pub struct ClientSessionBuilder {
81    client_addr: ClientAddress,
82    buffer_pool: BufferPool,
83    router: Option<Arc<BackendSelector>>,
84    routing_mode: RoutingMode,
85    auth_handler: Arc<AuthHandler>,
86    metrics: MetricsCollector,
87    connection_stats: Option<crate::metrics::ConnectionStatsAggregator>,
88    cache: BuilderCache,
89    cache_articles: bool,
90    adaptive_precheck: bool,
91}
92
93enum BuilderCache {
94    DefaultAvailability,
95    Shared(Arc<crate::cache::UnifiedCache>),
96}
97
98impl BuilderCache {
99    fn into_cache(self) -> Arc<crate::cache::UnifiedCache> {
100        match self {
101            Self::DefaultAvailability => ClientSession::default_cache(),
102            Self::Shared(cache) => cache,
103        }
104    }
105}
106
107impl ClientSessionBuilder {
108    /// Configure the session to use per-command routing with a backend router
109    ///
110    /// When a router is provided, the session will route each command independently
111    /// to potentially different backend servers.
112    #[must_use]
113    pub fn with_router(mut self, router: Arc<BackendSelector>) -> Self {
114        self.router = Some(router);
115        self
116    }
117
118    /// Set the routing mode for this session
119    ///
120    /// # Arguments
121    /// * `mode` - The routing mode (Stateful, `PerCommand`, or Hybrid)
122    ///
123    /// Note: If you use `with_router()`, you typically want `PerCommand` or Hybrid mode.
124    #[must_use]
125    pub const fn with_routing_mode(mut self, mode: RoutingMode) -> Self {
126        self.routing_mode = mode;
127        self
128    }
129
130    /// Set the authentication handler
131    #[must_use]
132    pub fn with_auth_handler(mut self, auth_handler: Arc<AuthHandler>) -> Self {
133        self.auth_handler = auth_handler;
134        self
135    }
136
137    /// Add connection stats aggregation to this session
138    #[must_use]
139    pub fn with_connection_stats(
140        mut self,
141        connection_stats: crate::metrics::ConnectionStatsAggregator,
142    ) -> Self {
143        self.connection_stats = Some(connection_stats);
144        self
145    }
146
147    /// Add article cache to this session (always present for backend availability tracking)
148    #[must_use]
149    pub fn with_cache(mut self, cache: Arc<crate::cache::UnifiedCache>) -> Self {
150        self.cache = BuilderCache::Shared(cache);
151        self
152    }
153
154    /// Set whether to cache article bodies.
155    ///
156    /// When false, only backend availability is tracked (saves memory).
157    /// When true, full article bodies are cached.
158    #[must_use]
159    pub const fn with_cache_articles(mut self, cache: bool) -> Self {
160        self.cache_articles = cache;
161        self
162    }
163
164    /// Configure adaptive availability prechecking
165    #[must_use]
166    pub const fn with_adaptive_precheck(mut self, enable: bool) -> Self {
167        self.adaptive_precheck = enable;
168        self
169    }
170
171    /// Build the client session
172    ///
173    /// Creates a new `ClientSession` with a unique client ID and the configured
174    /// routing mode.
175    #[must_use]
176    pub fn build(self) -> ClientSession {
177        let (mode, routing_mode) = match (&self.router, self.routing_mode) {
178            // Per-command or Hybrid: start in per-command mode (stateless)
179            (Some(_), RoutingMode::PerCommand | RoutingMode::Hybrid) => {
180                (SessionMode::PerCommand, self.routing_mode)
181            }
182            // Stateful mode with router, or no router: always Stateful
183            (Some(_), RoutingMode::Stateful) | (None, _) => {
184                (SessionMode::Stateful, RoutingMode::Stateful)
185            }
186        };
187
188        ClientSession {
189            client_addr: self.client_addr,
190            buffer_pool: self.buffer_pool,
191            client_id: ClientId::new(),
192            router: self.router,
193            mode_state: ModeState::new(mode, routing_mode),
194            auth_handler: self.auth_handler,
195            auth_state: AuthState::new(),
196            metrics: self.metrics,
197            connection_stats: self.connection_stats,
198            cache: self.cache.into_cache(),
199            cache_articles: self.cache_articles,
200            adaptive_precheck: self.adaptive_precheck,
201        }
202    }
203}
204
205impl ClientSession {
206    /// Create default cache for availability tracking only (no content caching)
207    fn default_cache() -> Arc<crate::cache::UnifiedCache> {
208        Arc::new(crate::cache::UnifiedCache::availability(
209            std::time::Duration::MAX,
210        ))
211    }
212
213    /// Create a new client session for 1:1 backend mapping
214    #[must_use]
215    pub fn new(
216        client_addr: ClientAddress,
217        buffer_pool: BufferPool,
218        auth_handler: Arc<AuthHandler>,
219        metrics: MetricsCollector,
220    ) -> Self {
221        Self {
222            client_addr,
223            buffer_pool,
224            client_id: ClientId::new(),
225            router: None,
226            mode_state: ModeState::new(SessionMode::Stateful, RoutingMode::Stateful),
227            auth_handler,
228            auth_state: AuthState::new(),
229            metrics,
230            connection_stats: None,
231            cache: Self::default_cache(),
232            cache_articles: false,
233            adaptive_precheck: false,
234        }
235    }
236
237    /// Create a new client session with per-command routing
238    ///
239    /// Each command will be routed to a potentially different backend server
240    /// using round-robin load balancing.
241    #[must_use]
242    pub fn new_with_router(
243        client_addr: ClientAddress,
244        buffer_pool: BufferPool,
245        router: Arc<BackendSelector>,
246        routing_mode: RoutingMode,
247        auth_handler: Arc<AuthHandler>,
248        metrics: MetricsCollector,
249    ) -> Self {
250        Self {
251            client_addr,
252            buffer_pool,
253            client_id: ClientId::new(),
254            router: Some(router),
255            mode_state: ModeState::new(SessionMode::PerCommand, routing_mode),
256            auth_handler,
257            auth_state: AuthState::new(),
258            metrics,
259            connection_stats: None,
260            cache: Self::default_cache(),
261            cache_articles: false,
262            adaptive_precheck: false,
263        }
264    }
265
266    /// Create a builder for constructing a client session
267    ///
268    /// # Examples
269    ///
270    /// ```
271    /// use std::net::SocketAddr;
272    /// use std::sync::Arc;
273    /// use nntp_proxy::session::ClientSession;
274    /// use nntp_proxy::pool::BufferPool;
275    /// use nntp_proxy::types::BufferSize;
276    /// use nntp_proxy::auth::AuthHandler;
277    /// use nntp_proxy::metrics::MetricsCollector;
278    ///
279    /// let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
280    /// let buffer_pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 10);
281    /// let auth_handler = Arc::new(AuthHandler::new(None, None).unwrap());
282    /// let metrics = MetricsCollector::new(1);
283    ///
284    /// let session = ClientSession::builder(addr.into(), buffer_pool, auth_handler, metrics)
285    ///     .build();
286    /// ```
287    #[must_use]
288    pub fn builder(
289        client_addr: ClientAddress,
290        buffer_pool: BufferPool,
291        auth_handler: Arc<AuthHandler>,
292        metrics: MetricsCollector,
293    ) -> ClientSessionBuilder {
294        ClientSessionBuilder {
295            client_addr,
296            buffer_pool,
297            router: None,
298            routing_mode: RoutingMode::Stateful,
299            auth_handler,
300            metrics,
301            connection_stats: None,
302            cache: BuilderCache::DefaultAvailability,
303            cache_articles: false,
304            adaptive_precheck: false,
305        }
306    }
307
308    // Getters and helper methods
309
310    /// Get the unique client ID
311    #[must_use]
312    #[inline]
313    pub const fn client_id(&self) -> ClientId {
314        self.client_id
315    }
316
317    /// Check if this session is using per-command routing
318    ///
319    /// Returns true if this session has a router available (regardless of current mode).
320    /// This is slightly different from checking routing mode - a session can have a router
321    /// but be in Stateful mode (e.g., after hybrid mode switches).
322    #[must_use]
323    #[inline]
324    pub const fn is_per_command_routing(&self) -> bool {
325        self.router.is_some()
326    }
327
328    /// Get the current session mode
329    #[must_use]
330    #[inline]
331    pub fn mode(&self) -> SessionMode {
332        self.mode_state.mode()
333    }
334
335    /// Get the authenticated username (if any) - zero-cost reference
336    ///
337    /// Returns the authenticated username as an `Arc<str>` for cheap cloning.
338    /// Returns None if the client has not authenticated yet.
339    #[inline]
340    #[must_use]
341    pub fn username(&self) -> Option<&str> {
342        self.auth_state.username()
343    }
344
345    /// Set the authenticated username (write-once)
346    ///
347    /// This marks the session as authenticated and stores the username.
348    /// Called after successful authentication with the backend.
349    pub(crate) fn set_username(&self, username: Option<String>) {
350        if let Some(name) = username {
351            self.auth_state.mark_authenticated(name);
352        }
353    }
354
355    /// Get the connection stats aggregator (if enabled)
356    #[must_use]
357    #[inline]
358    pub(crate) const fn connection_stats(
359        &self,
360    ) -> Option<&crate::metrics::ConnectionStatsAggregator> {
361        self.connection_stats.as_ref()
362    }
363
364    /// Check if already authenticated (cached for performance)
365    ///
366    /// # Arguments
367    /// * `skip_auth_check` - If true, bypasses the authentication check
368    ///
369    /// # Returns
370    /// Returns true if authenticated or if `skip_auth_check` is true
371    #[inline]
372    pub(crate) fn is_authenticated_cached(&self, skip_auth_check: bool) -> bool {
373        self.auth_state.is_authenticated_or_skipped(skip_auth_check)
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380    use crate::auth::AuthHandler;
381    use crate::metrics::MetricsCollector;
382    use crate::types::BufferSize;
383    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
384    use std::sync::Arc;
385
386    /// Helper to create a default `AuthHandler` for tests (no auth)
387    fn test_auth_handler() -> Arc<AuthHandler> {
388        Arc::new(AuthHandler::new(None, None).unwrap())
389    }
390
391    /// Helper to create a `MetricsCollector` for tests
392    fn test_metrics() -> MetricsCollector {
393        MetricsCollector::new(1)
394    }
395
396    #[test]
397    fn test_client_session_creation() {
398        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
399        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
400        let session = ClientSession::new(
401            addr.into(),
402            buffer_pool,
403            test_auth_handler(),
404            test_metrics(),
405        );
406
407        assert_eq!(session.client_addr.port(), 8080);
408        assert_eq!(session.client_addr.ip(), IpAddr::V4(Ipv4Addr::LOCALHOST));
409    }
410
411    #[test]
412    fn test_client_session_with_different_ports() {
413        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
414
415        let addr1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
416        let session1 = ClientSession::new(
417            addr1.into(),
418            buffer_pool.clone(),
419            test_auth_handler(),
420            test_metrics(),
421        );
422
423        let addr2 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9090);
424        let session2 = ClientSession::new(
425            addr2.into(),
426            buffer_pool,
427            test_auth_handler(),
428            test_metrics(),
429        );
430
431        assert_ne!(session1.client_addr.port(), session2.client_addr.port());
432        assert_eq!(session1.client_addr.port(), 8080);
433        assert_eq!(session2.client_addr.port(), 9090);
434    }
435
436    #[test]
437    fn test_client_session_with_ipv6() {
438        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
439        let addr = SocketAddr::new(IpAddr::V6("::1".parse().unwrap()), 8119);
440        let session = ClientSession::new(
441            addr.into(),
442            buffer_pool,
443            test_auth_handler(),
444            test_metrics(),
445        );
446
447        assert_eq!(session.client_addr.port(), 8119);
448        assert!(session.client_addr.is_ipv6());
449    }
450
451    #[test]
452    fn test_buffer_pool_cloning() {
453        let buffer_pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 10);
454        let buffer_pool_clone = buffer_pool.clone();
455
456        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), 1234);
457        let _session1 = ClientSession::new(
458            addr.into(),
459            buffer_pool,
460            test_auth_handler(),
461            test_metrics(),
462        );
463        let _session2 = ClientSession::new(
464            addr.into(),
465            buffer_pool_clone,
466            test_auth_handler(),
467            test_metrics(),
468        );
469    }
470
471    #[test]
472    fn test_session_addr_formatting() {
473        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 5555);
474        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
475        let session = ClientSession::new(
476            addr.into(),
477            buffer_pool,
478            test_auth_handler(),
479            test_metrics(),
480        );
481
482        let addr_str = format!("{}", session.client_addr);
483        assert!(addr_str.contains("10.0.0.1"));
484        assert!(addr_str.contains("5555"));
485    }
486
487    #[test]
488    fn test_multiple_sessions_same_buffer_pool() {
489        let buffer_pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 8);
490        let sessions: Vec<_> = (0..5)
491            .map(|i| {
492                let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8000 + i);
493                ClientSession::new(
494                    addr.into(),
495                    buffer_pool.clone(),
496                    test_auth_handler(),
497                    test_metrics(),
498                )
499            })
500            .collect();
501
502        assert_eq!(sessions.len(), 5);
503        for (i, session) in sessions.iter().enumerate() {
504            assert_eq!(session.client_addr.port(), 8000 + i as u16);
505        }
506    }
507
508    #[test]
509    fn test_loopback_address() {
510        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
511        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8119);
512        let session = ClientSession::new(
513            addr.into(),
514            buffer_pool,
515            test_auth_handler(),
516            test_metrics(),
517        );
518
519        assert!(session.client_addr.ip().is_loopback());
520    }
521
522    #[test]
523    fn test_unspecified_address() {
524        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
525        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0);
526        let session = ClientSession::new(
527            addr.into(),
528            buffer_pool,
529            test_auth_handler(),
530            test_metrics(),
531        );
532
533        assert!(session.client_addr.ip().is_unspecified());
534        assert_eq!(session.client_addr.port(), 0);
535    }
536
537    #[test]
538    fn test_session_without_router() {
539        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
540        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
541        let session = ClientSession::new(
542            addr.into(),
543            buffer_pool,
544            test_auth_handler(),
545            test_metrics(),
546        );
547
548        assert!(!session.is_per_command_routing());
549        assert_eq!(session.client_addr.port(), 8080);
550    }
551
552    #[test]
553    fn test_session_with_router() {
554        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
555        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
556        let router = Arc::new(BackendSelector::new());
557        let session = ClientSession::new_with_router(
558            addr.into(),
559            buffer_pool,
560            router,
561            RoutingMode::PerCommand,
562            test_auth_handler(),
563            test_metrics(),
564        );
565
566        assert!(session.is_per_command_routing());
567        assert_eq!(session.client_addr.port(), 8080);
568    }
569
570    #[test]
571    fn test_client_id_uniqueness() {
572        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
573        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
574
575        let session1 = ClientSession::new(
576            addr.into(),
577            buffer_pool.clone(),
578            test_auth_handler(),
579            test_metrics(),
580        );
581        let session2 = ClientSession::new(
582            addr.into(),
583            buffer_pool,
584            test_auth_handler(),
585            test_metrics(),
586        );
587
588        assert_ne!(session1.client_id(), session2.client_id());
589    }
590
591    #[test]
592    fn test_session_mode_enum() {
593        assert_eq!(SessionMode::PerCommand, SessionMode::PerCommand);
594        assert_eq!(SessionMode::Stateful, SessionMode::Stateful);
595        assert_ne!(SessionMode::PerCommand, SessionMode::Stateful);
596
597        let per_command = format!("{:?}", SessionMode::PerCommand);
598        let stateful = format!("{:?}", SessionMode::Stateful);
599        assert!(per_command.contains("PerCommand"));
600        assert!(stateful.contains("Stateful"));
601    }
602
603    #[test]
604    fn test_hybrid_session_creation() {
605        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
606        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
607        let router = Arc::new(BackendSelector::new());
608
609        let session = ClientSession::new_with_router(
610            addr.into(),
611            buffer_pool,
612            router,
613            RoutingMode::Hybrid,
614            test_auth_handler(),
615            test_metrics(),
616        );
617
618        assert!(session.is_per_command_routing());
619        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
620        assert_eq!(session.mode(), SessionMode::PerCommand);
621    }
622
623    #[test]
624    fn test_routing_mode_configurations() {
625        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
626        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
627        let router = Arc::new(BackendSelector::new());
628
629        // Stateful mode
630        let session = ClientSession::new_with_router(
631            addr.into(),
632            buffer_pool.clone(),
633            router.clone(),
634            RoutingMode::Stateful,
635            test_auth_handler(),
636            test_metrics(),
637        );
638        assert!(session.is_per_command_routing());
639        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
640
641        // PerCommand mode
642        let session = ClientSession::new_with_router(
643            addr.into(),
644            buffer_pool.clone(),
645            router.clone(),
646            RoutingMode::PerCommand,
647            test_auth_handler(),
648            test_metrics(),
649        );
650        assert!(session.is_per_command_routing());
651        assert_eq!(session.mode_state.routing_mode(), RoutingMode::PerCommand);
652        assert_eq!(session.mode(), SessionMode::PerCommand);
653
654        // Hybrid mode
655        let session = ClientSession::new_with_router(
656            addr.into(),
657            buffer_pool,
658            router,
659            RoutingMode::Hybrid,
660            test_auth_handler(),
661            test_metrics(),
662        );
663        assert!(session.is_per_command_routing());
664        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
665        assert_eq!(session.mode(), SessionMode::PerCommand);
666    }
667
668    #[test]
669    fn test_hybrid_mode_initial_state() {
670        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
671        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
672        let router = Arc::new(BackendSelector::new());
673
674        let session = ClientSession::new_with_router(
675            addr.into(),
676            buffer_pool,
677            router,
678            RoutingMode::Hybrid,
679            test_auth_handler(),
680            test_metrics(),
681        );
682
683        assert_eq!(session.mode(), SessionMode::PerCommand);
684        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
685        assert!(session.is_per_command_routing());
686    }
687
688    #[test]
689    fn test_is_per_command_routing_logic() {
690        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
691        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
692        let router = Arc::new(BackendSelector::new());
693
694        // Stateful mode has router capability
695        let session = ClientSession::new_with_router(
696            addr.into(),
697            buffer_pool.clone(),
698            router.clone(),
699            RoutingMode::Stateful,
700            test_auth_handler(),
701            test_metrics(),
702        );
703        assert!(session.is_per_command_routing());
704
705        // PerCommand mode
706        let session = ClientSession::new_with_router(
707            addr.into(),
708            buffer_pool.clone(),
709            router.clone(),
710            RoutingMode::PerCommand,
711            test_auth_handler(),
712            test_metrics(),
713        );
714        assert!(session.is_per_command_routing());
715
716        // Hybrid mode (initially)
717        let session = ClientSession::new_with_router(
718            addr.into(),
719            buffer_pool.clone(),
720            router,
721            RoutingMode::Hybrid,
722            test_auth_handler(),
723            test_metrics(),
724        );
725        assert!(session.is_per_command_routing());
726
727        // Session without router
728        let session = ClientSession::new(
729            addr.into(),
730            buffer_pool,
731            test_auth_handler(),
732            test_metrics(),
733        );
734        assert!(!session.is_per_command_routing());
735    }
736
737    // ==================== ClientSessionBuilder Tests ====================
738
739    #[test]
740    fn test_builder_basic_construction() {
741        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
742        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
743        let auth_handler = test_auth_handler();
744
745        let session =
746            ClientSession::builder(addr.into(), buffer_pool, auth_handler, test_metrics()).build();
747
748        assert_eq!(*session.client_addr, addr);
749        assert!(!session.is_per_command_routing());
750        assert_eq!(session.mode(), SessionMode::Stateful);
751        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
752    }
753
754    #[test]
755    fn test_builder_with_router() {
756        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
757        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
758        let router = Arc::new(BackendSelector::new());
759
760        let session = ClientSession::builder(
761            addr.into(),
762            buffer_pool,
763            test_auth_handler(),
764            test_metrics(),
765        )
766        .with_router(router)
767        .with_routing_mode(RoutingMode::PerCommand)
768        .build();
769
770        assert!(session.is_per_command_routing());
771        assert_eq!(session.mode(), SessionMode::PerCommand);
772        assert_eq!(session.mode_state.routing_mode(), RoutingMode::PerCommand);
773    }
774
775    #[test]
776    fn test_builder_with_hybrid_mode() {
777        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
778        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
779        let router = Arc::new(BackendSelector::new());
780
781        let session = ClientSession::builder(
782            addr.into(),
783            buffer_pool,
784            test_auth_handler(),
785            test_metrics(),
786        )
787        .with_router(router)
788        .with_routing_mode(RoutingMode::Hybrid)
789        .build();
790
791        assert!(session.is_per_command_routing());
792        assert_eq!(session.mode(), SessionMode::PerCommand);
793        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
794    }
795
796    #[test]
797    fn test_builder_with_stateful_mode() {
798        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
799        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
800        let router = Arc::new(BackendSelector::new());
801
802        // Builder with router but Stateful mode requested
803        let session = ClientSession::builder(
804            addr.into(),
805            buffer_pool,
806            test_auth_handler(),
807            test_metrics(),
808        )
809        .with_router(router)
810        .with_routing_mode(RoutingMode::Stateful)
811        .build();
812
813        assert!(session.is_per_command_routing());
814        assert_eq!(session.mode(), SessionMode::Stateful);
815        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
816    }
817
818    #[test]
819    fn test_builder_without_router_ignores_mode() {
820        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
821        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
822
823        // Request PerCommand mode but no router - should default to Stateful
824        let session = ClientSession::builder(
825            addr.into(),
826            buffer_pool,
827            test_auth_handler(),
828            test_metrics(),
829        )
830        .with_routing_mode(RoutingMode::PerCommand)
831        .build();
832
833        assert!(!session.is_per_command_routing());
834        assert_eq!(session.mode(), SessionMode::Stateful);
835        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
836    }
837
838    #[test]
839    fn test_builder_with_auth_handler() {
840        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
841        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
842        let auth_handler =
843            Arc::new(AuthHandler::new(Some("user".to_string()), Some("pass".to_string())).unwrap());
844
845        let session = ClientSession::builder(
846            addr.into(),
847            buffer_pool,
848            test_auth_handler(),
849            test_metrics(),
850        )
851        .with_auth_handler(auth_handler)
852        .build();
853
854        // Verify auth_handler is set (can't test internals, but creation succeeds)
855        assert_eq!(*session.client_addr, addr);
856    }
857
858    #[test]
859    fn test_builder_with_metrics() {
860        use crate::metrics::MetricsCollector;
861
862        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
863        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
864        let metrics = MetricsCollector::new(1); // 1 backend
865
866        let session =
867            ClientSession::builder(addr.into(), buffer_pool, test_auth_handler(), metrics).build();
868
869        // Metrics is always present now, not optional
870        let _ = session.metrics; // Just verify it exists
871    }
872
873    #[test]
874    fn test_builder_with_connection_stats() {
875        use crate::metrics::ConnectionStatsAggregator;
876
877        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
878        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
879        let stats = ConnectionStatsAggregator::default();
880
881        let session = ClientSession::builder(
882            addr.into(),
883            buffer_pool,
884            test_auth_handler(),
885            test_metrics(),
886        )
887        .with_connection_stats(stats)
888        .build();
889
890        assert!(session.connection_stats().is_some());
891    }
892
893    #[test]
894    fn test_builder_with_cache() {
895        use crate::cache::UnifiedCache;
896        use std::time::Duration;
897
898        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
899        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
900        let cache = Arc::new(UnifiedCache::memory(100, Duration::from_secs(3600)));
901
902        let session = ClientSession::builder(
903            addr.into(),
904            buffer_pool,
905            test_auth_handler(),
906            test_metrics(),
907        )
908        .with_cache(cache)
909        .build();
910
911        // Cache is always present now (Arc not Option)
912        assert_eq!(session.cache.capacity(), 100);
913    }
914
915    #[test]
916    fn test_builder_method_chaining() {
917        use crate::metrics::MetricsCollector;
918
919        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
920        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
921        let router = Arc::new(BackendSelector::new());
922        let metrics = MetricsCollector::new(1); // 1 backend
923
924        // Chain all builder methods
925        let session =
926            ClientSession::builder(addr.into(), buffer_pool, test_auth_handler(), metrics)
927                .with_router(router)
928                .with_routing_mode(RoutingMode::Hybrid)
929                .build();
930
931        assert!(session.is_per_command_routing());
932        assert_eq!(session.mode(), SessionMode::PerCommand);
933        assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
934        // Metrics is always present now, not optional
935        let _ = session.metrics; // Just verify it exists
936    }
937
938    // ==================== Session Business Logic Tests ====================
939
940    #[test]
941    fn test_mode_getter() {
942        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
943        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
944        let router = Arc::new(BackendSelector::new());
945
946        // Per-command mode
947        let session = ClientSession::new_with_router(
948            addr.into(),
949            buffer_pool.clone(),
950            router,
951            RoutingMode::PerCommand,
952            test_auth_handler(),
953            test_metrics(),
954        );
955        assert_eq!(session.mode(), SessionMode::PerCommand);
956
957        // Stateful mode (no router)
958        let session = ClientSession::new(
959            addr.into(),
960            buffer_pool,
961            test_auth_handler(),
962            test_metrics(),
963        );
964        assert_eq!(session.mode(), SessionMode::Stateful);
965    }
966
967    #[test]
968    fn test_username_initially_none() {
969        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
970        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
971        let session = ClientSession::new(
972            addr.into(),
973            buffer_pool,
974            test_auth_handler(),
975            test_metrics(),
976        );
977
978        assert!(session.username().is_none());
979    }
980
981    #[test]
982    fn test_set_username_and_get() {
983        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
984        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
985        let session = ClientSession::new(
986            addr.into(),
987            buffer_pool,
988            test_auth_handler(),
989            test_metrics(),
990        );
991
992        session.set_username(Some("testuser".to_string()));
993
994        let username = session.username();
995        assert!(username.is_some());
996        assert_eq!(username, Some("testuser"));
997    }
998
999    #[test]
1000    fn test_set_username_none() {
1001        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1002        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1003        let session = ClientSession::new(
1004            addr.into(),
1005            buffer_pool,
1006            test_auth_handler(),
1007            test_metrics(),
1008        );
1009
1010        session.set_username(None);
1011
1012        assert!(session.username().is_none());
1013    }
1014
1015    #[test]
1016    fn test_username_cheap_clone() {
1017        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1018        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1019        let session = ClientSession::new(
1020            addr.into(),
1021            buffer_pool,
1022            test_auth_handler(),
1023            test_metrics(),
1024        );
1025
1026        session.set_username(Some("testuser".to_string()));
1027
1028        let username1 = session.username();
1029        let username2 = session.username();
1030
1031        assert!(username1.is_some());
1032        assert!(username2.is_some());
1033        assert_eq!(username1, username2);
1034        assert_eq!(username1, Some("testuser"));
1035    }
1036
1037    #[test]
1038    fn test_connection_stats_none_by_default() {
1039        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1040        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1041        let session = ClientSession::new(
1042            addr.into(),
1043            buffer_pool,
1044            test_auth_handler(),
1045            test_metrics(),
1046        );
1047
1048        assert!(session.connection_stats().is_none());
1049    }
1050
1051    #[test]
1052    fn test_connection_stats_with_aggregator() {
1053        use crate::metrics::ConnectionStatsAggregator;
1054
1055        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1056        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1057        let stats = ConnectionStatsAggregator::default();
1058
1059        let session = ClientSession::builder(
1060            addr.into(),
1061            buffer_pool,
1062            test_auth_handler(),
1063            test_metrics(),
1064        )
1065        .with_connection_stats(stats)
1066        .build();
1067
1068        assert!(session.connection_stats().is_some());
1069    }
1070
1071    // ==================== Metrics Helper Methods Tests ====================
1072
1073    #[test]
1074    fn test_metrics_direct_calls_no_metrics() {
1075        use crate::types::BackendId;
1076
1077        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1078        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1079        let session = ClientSession::new(
1080            addr.into(),
1081            buffer_pool,
1082            test_auth_handler(),
1083            test_metrics(),
1084        );
1085
1086        // Metrics is always present now
1087        session.metrics.record_command(BackendId::from_index(0));
1088        session.metrics.user_command(session.username());
1089        session.metrics.stateful_session_started();
1090        session.metrics.stateful_session_ended();
1091        session.metrics.user_bytes_sent(session.username(), 1024);
1092        session
1093            .metrics
1094            .user_bytes_received(session.username(), 2048);
1095    }
1096
1097    #[test]
1098    fn test_metrics_direct_calls_with_metrics() {
1099        use crate::metrics::MetricsCollector;
1100        use crate::types::BackendId;
1101
1102        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1103        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1104        let metrics = MetricsCollector::new(1); // 1 backend
1105
1106        let session = ClientSession::builder(
1107            addr.into(),
1108            buffer_pool,
1109            test_auth_handler(),
1110            metrics.clone(),
1111        )
1112        .build();
1113
1114        session.set_username(Some("testuser".to_string()));
1115
1116        // Call metrics directly - should record
1117        session.metrics.record_command(BackendId::from_index(0));
1118        session.metrics.user_command(session.username());
1119        session.metrics.stateful_session_started();
1120        session.metrics.stateful_session_ended();
1121        session.metrics.user_bytes_sent(session.username(), 1024);
1122        session
1123            .metrics
1124            .user_bytes_received(session.username(), 2048);
1125
1126        // Verify metrics were recorded (snapshot should have data)
1127        let snapshot = metrics.snapshot(None);
1128        assert!(!snapshot.backend_stats.is_empty());
1129        assert!(snapshot.backend_stats[0].total_commands.get() > 0);
1130    }
1131
1132    #[test]
1133    fn test_record_command_with_metrics() {
1134        use crate::metrics::MetricsCollector;
1135        use crate::types::BackendId;
1136
1137        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1138        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1139        let metrics = MetricsCollector::new(1); // 1 backend
1140
1141        let session = ClientSession::builder(
1142            addr.into(),
1143            buffer_pool,
1144            test_auth_handler(),
1145            metrics.clone(),
1146        )
1147        .build();
1148
1149        let backend_id = BackendId::from_index(0);
1150        session.metrics.record_command(backend_id);
1151
1152        let snapshot = metrics.snapshot(None);
1153        assert_eq!(snapshot.backend_stats[0].total_commands.get(), 1);
1154    }
1155
1156    #[test]
1157    fn test_user_bytes_tracking() {
1158        use crate::metrics::MetricsCollector;
1159
1160        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1161        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1162        let metrics = MetricsCollector::new(1); // 1 backend
1163
1164        let session = ClientSession::builder(
1165            addr.into(),
1166            buffer_pool,
1167            test_auth_handler(),
1168            metrics.clone(),
1169        )
1170        .build();
1171
1172        session.set_username(Some("testuser".to_string()));
1173        session.metrics.user_bytes_sent(session.username(), 1024);
1174        session
1175            .metrics
1176            .user_bytes_received(session.username(), 2048);
1177
1178        let snapshot = metrics.snapshot(None);
1179        let user_stats = snapshot
1180            .user_stats
1181            .iter()
1182            .find(|s| s.username == "testuser");
1183        assert!(user_stats.is_some());
1184
1185        let stats = user_stats.unwrap();
1186        assert_eq!(stats.bytes_sent.as_u64(), 1024);
1187        assert_eq!(stats.bytes_received.as_u64(), 2048);
1188    }
1189
1190    #[test]
1191    fn test_stateful_session_tracking() {
1192        use crate::metrics::MetricsCollector;
1193
1194        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1195        let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1196        let metrics = MetricsCollector::new(1); // 1 backend
1197
1198        let session = ClientSession::builder(
1199            addr.into(),
1200            buffer_pool,
1201            test_auth_handler(),
1202            metrics.clone(),
1203        )
1204        .build();
1205
1206        session.metrics.stateful_session_started();
1207
1208        let snapshot = metrics.snapshot(None);
1209        assert_eq!(snapshot.stateful_sessions, 1);
1210
1211        session.metrics.stateful_session_ended();
1212
1213        let snapshot = metrics.snapshot(None);
1214        assert_eq!(snapshot.stateful_sessions, 0);
1215    }
1216}