gemini_genai_rs/transport/
builder.rs1use 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
10pub struct ConnectBuilder<T = TungsteniteTransport, C = JsonCodec> {
29 config: SessionConfig,
30 transport_config: TransportConfig,
31 transport: T,
32 codec: C,
33}
34
35impl ConnectBuilder {
36 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 pub fn transport_config(mut self, tc: TransportConfig) -> Self {
50 self.transport_config = tc;
51 self
52 }
53
54 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 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 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}