gemini_genai_rs/transport/
builder.rs

1//! ConnectBuilder — ergonomic builder for advanced transport/codec configuration.
2
3use crate::protocol::types::SessionConfig;
4use crate::session::{SessionError, SessionHandle};
5use crate::transport::codec::{Codec, JsonCodec};
6use crate::transport::connection::connect_with;
7use crate::transport::ws::{Transport, TungsteniteTransport};
8use crate::transport::TransportConfig;
9
10/// Builder for advanced connection configuration.
11///
12/// Allows customizing the transport and codec used for the connection.
13/// Defaults to TungsteniteTransport + JsonCodec.
14///
15/// # Example
16/// ```rust,no_run
17/// use gemini_genai_rs::prelude::*;
18///
19/// # async fn example() {
20/// let config = SessionConfig::new("key");
21/// let handle = ConnectBuilder::new(config)
22///     .transport_config(TransportConfig { connect_timeout_secs: 30, ..Default::default() })
23///     .build()
24///     .await
25///     .unwrap();
26/// # }
27/// ```
28pub struct ConnectBuilder<T = TungsteniteTransport, C = JsonCodec> {
29    config: SessionConfig,
30    transport_config: TransportConfig,
31    transport: T,
32    codec: C,
33}
34
35impl ConnectBuilder {
36    /// Create a new builder with default transport and codec.
37    pub fn new(config: SessionConfig) -> Self {
38        Self {
39            config,
40            transport_config: TransportConfig::default(),
41            transport: TungsteniteTransport::new(),
42            codec: JsonCodec,
43        }
44    }
45}
46
47impl<T: Transport, C: Codec> ConnectBuilder<T, C> {
48    /// Set the transport configuration.
49    pub fn transport_config(mut self, tc: TransportConfig) -> Self {
50        self.transport_config = tc;
51        self
52    }
53
54    /// Use a custom transport implementation.
55    pub fn transport<T2: Transport>(self, transport: T2) -> ConnectBuilder<T2, C> {
56        ConnectBuilder {
57            config: self.config,
58            transport_config: self.transport_config,
59            transport,
60            codec: self.codec,
61        }
62    }
63
64    /// Use a custom codec implementation.
65    pub fn codec<C2: Codec>(self, codec: C2) -> ConnectBuilder<T, C2> {
66        ConnectBuilder {
67            config: self.config,
68            transport_config: self.transport_config,
69            transport: self.transport,
70            codec,
71        }
72    }
73
74    /// Build the connection and return a SessionHandle.
75    pub async fn build(self) -> Result<SessionHandle, SessionError> {
76        connect_with(
77            self.config,
78            self.transport_config,
79            self.transport,
80            self.codec,
81        )
82        .await
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::protocol::types::*;
90    use crate::transport::ws::MockTransport;
91
92    #[test]
93    fn builder_compiles_with_defaults() {
94        let config = SessionConfig::new("key").model(GeminiModel::Gemini2_0FlashLive);
95        let _builder = ConnectBuilder::new(config);
96    }
97
98    #[test]
99    fn builder_with_custom_transport_config() {
100        let config = SessionConfig::new("key");
101        let _builder = ConnectBuilder::new(config).transport_config(TransportConfig {
102            connect_timeout_secs: 30,
103            ..Default::default()
104        });
105    }
106
107    #[test]
108    fn builder_with_mock_transport() {
109        let config = SessionConfig::new("key");
110        let mock = MockTransport::new();
111        let _builder = ConnectBuilder::new(config).transport(mock);
112    }
113
114    #[test]
115    fn builder_with_custom_codec() {
116        let config = SessionConfig::new("key");
117        let _builder = ConnectBuilder::new(config).codec(JsonCodec);
118    }
119
120    #[tokio::test]
121    async fn builder_with_mock_builds() {
122        let mut mock = MockTransport::new();
123        mock.script_recv(br#"{"setupComplete":{}}"#.to_vec());
124
125        let config = SessionConfig::new("key").model(GeminiModel::Gemini2_0FlashLive);
126        let handle = ConnectBuilder::new(config)
127            .transport(mock)
128            .build()
129            .await
130            .unwrap();
131
132        handle
133            .wait_for_phase(crate::session::SessionPhase::Active)
134            .await;
135        assert_eq!(handle.phase(), crate::session::SessionPhase::Active);
136    }
137}