1pub(crate) mod backends;
2
3use crate::client::NblmClient;
4use crate::error::Result;
5use crate::models::enterprise::{
6 audio::{AudioOverviewRequest, AudioOverviewResponse},
7 notebook::{
8 BatchDeleteNotebooksRequest, BatchDeleteNotebooksResponse, ListRecentlyViewedResponse,
9 Notebook,
10 },
11 share::{AccountRole, ShareResponse},
12 source::{
13 BatchCreateSourcesRequest, BatchCreateSourcesResponse, BatchDeleteSourcesRequest,
14 BatchDeleteSourcesResponse, NotebookSource, UploadSourceFileResponse, UserContent,
15 },
16};
17
18impl NblmClient {
19 pub async fn create_notebook(&self, title: impl Into<String>) -> Result<Notebook> {
20 self.backends
21 .notebooks()
22 .create_notebook(title.into())
23 .await
24 }
25
26 pub async fn batch_delete_notebooks(
27 &self,
28 request: BatchDeleteNotebooksRequest,
29 ) -> Result<BatchDeleteNotebooksResponse> {
30 self.backends
31 .notebooks()
32 .batch_delete_notebooks(request)
33 .await
34 }
35
36 pub async fn delete_notebooks(
37 &self,
38 notebook_names: Vec<String>,
39 ) -> Result<BatchDeleteNotebooksResponse> {
40 self.backends
41 .notebooks()
42 .delete_notebooks(notebook_names)
43 .await
44 }
45
46 pub async fn share_notebook(
47 &self,
48 notebook_id: &str,
49 accounts: Vec<AccountRole>,
50 ) -> Result<ShareResponse> {
51 self.backends
52 .notebooks()
53 .share_notebook(notebook_id, accounts)
54 .await
55 }
56
57 pub async fn list_recently_viewed(
58 &self,
59 page_size: Option<u32>,
60 ) -> Result<ListRecentlyViewedResponse> {
61 self.backends
62 .notebooks()
63 .list_recently_viewed(page_size)
64 .await
65 }
66
67 pub async fn batch_create_sources(
68 &self,
69 notebook_id: &str,
70 request: BatchCreateSourcesRequest,
71 ) -> Result<BatchCreateSourcesResponse> {
72 let includes_drive = has_drive_content(request.user_contents.iter());
73 self.ensure_drive_scope_if_needed(includes_drive).await?;
74 self.backends
75 .sources()
76 .batch_create_sources(notebook_id, request)
77 .await
78 }
79
80 pub async fn add_sources(
81 &self,
82 notebook_id: &str,
83 contents: Vec<UserContent>,
84 ) -> Result<BatchCreateSourcesResponse> {
85 let includes_drive = has_drive_content(contents.iter());
86 self.ensure_drive_scope_if_needed(includes_drive).await?;
87 self.backends
88 .sources()
89 .add_sources(notebook_id, contents)
90 .await
91 }
92
93 pub async fn batch_delete_sources(
94 &self,
95 notebook_id: &str,
96 request: BatchDeleteSourcesRequest,
97 ) -> Result<BatchDeleteSourcesResponse> {
98 self.backends
99 .sources()
100 .batch_delete_sources(notebook_id, request)
101 .await
102 }
103
104 pub async fn delete_sources(
105 &self,
106 notebook_id: &str,
107 source_names: Vec<String>,
108 ) -> Result<BatchDeleteSourcesResponse> {
109 self.backends
110 .sources()
111 .delete_sources(notebook_id, source_names)
112 .await
113 }
114
115 pub async fn upload_source_file(
116 &self,
117 notebook_id: &str,
118 file_name: &str,
119 content_type: &str,
120 data: Vec<u8>,
121 ) -> Result<UploadSourceFileResponse> {
122 self.backends
123 .sources()
124 .upload_source_file(notebook_id, file_name, content_type, data)
125 .await
126 }
127
128 pub async fn get_source(&self, notebook_id: &str, source_id: &str) -> Result<NotebookSource> {
129 self.backends
130 .sources()
131 .get_source(notebook_id, source_id)
132 .await
133 }
134
135 pub async fn create_audio_overview(
136 &self,
137 notebook_id: &str,
138 request: AudioOverviewRequest,
139 ) -> Result<AudioOverviewResponse> {
140 self.backends
141 .audio()
142 .create_audio_overview(notebook_id, request)
143 .await
144 }
145
146 pub async fn delete_audio_overview(&self, notebook_id: &str) -> Result<()> {
147 self.backends
148 .audio()
149 .delete_audio_overview(notebook_id)
150 .await
151 }
152}
153
154fn has_drive_content<'a, I>(contents: I) -> bool
155where
156 I: IntoIterator<Item = &'a UserContent>,
157{
158 contents
159 .into_iter()
160 .any(|content| matches!(content, UserContent::GoogleDrive { .. }))
161}
162
163#[cfg(test)]
164mod tests {
165 use super::*;
166 use crate::auth::StaticTokenProvider;
167 use crate::env::EnvironmentConfig;
168 use crate::error::Error;
169 use serde_json::json;
170 use serial_test::serial;
171 use std::sync::Arc;
172 use wiremock::matchers::{method, path, query_param};
173 use wiremock::{Mock, MockServer, ResponseTemplate};
174
175 struct EnvGuard {
176 key: &'static str,
177 original: Option<String>,
178 }
179
180 impl EnvGuard {
181 fn new(key: &'static str) -> Self {
182 let original = std::env::var(key).ok();
183 Self { key, original }
184 }
185 }
186
187 impl Drop for EnvGuard {
188 fn drop(&mut self) {
189 if let Some(value) = &self.original {
190 std::env::set_var(self.key, value);
191 } else {
192 std::env::remove_var(self.key);
193 }
194 }
195 }
196
197 async fn build_client(base_url: &str) -> NblmClient {
198 let provider = Arc::new(StaticTokenProvider::new("test-token"));
199 let env = EnvironmentConfig::enterprise("123", "global", "us").unwrap();
200 NblmClient::new(provider, env)
201 .unwrap()
202 .with_base_url(base_url)
203 .unwrap()
204 }
205
206 #[tokio::test]
207 #[serial]
208 async fn add_sources_validates_drive_scope() {
209 let server = MockServer::start().await;
210 let tokeninfo_url = format!("{}/tokeninfo", server.uri());
211 let _guard = EnvGuard::new("NBLM_TOKENINFO_ENDPOINT");
212 std::env::set_var("NBLM_TOKENINFO_ENDPOINT", &tokeninfo_url);
213
214 Mock::given(method("GET"))
215 .and(path("/tokeninfo"))
216 .and(query_param("access_token", "test-token"))
217 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
218 "scope": "https://www.googleapis.com/auth/drive.file"
219 })))
220 .expect(1)
221 .mount(&server)
222 .await;
223
224 Mock::given(method("POST"))
225 .and(path(
226 "/v1alpha/projects/123/locations/global/notebooks/notebook-id/sources:batchCreate",
227 ))
228 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
229 "sources": [],
230 "errorCount": 0
231 })))
232 .expect(1)
233 .mount(&server)
234 .await;
235
236 let client = build_client(&format!("{}/v1alpha", server.uri())).await;
237
238 let result = client
239 .add_sources(
240 "notebook-id",
241 vec![UserContent::google_drive(
242 "doc".to_string(),
243 "application/pdf".to_string(),
244 None,
245 )],
246 )
247 .await;
248
249 assert!(
250 result.is_ok(),
251 "expected add_sources to succeed: {:?}",
252 result
253 );
254 }
255
256 #[tokio::test]
257 #[serial]
258 async fn add_sources_rejects_missing_drive_scope() {
259 let server = MockServer::start().await;
260 let tokeninfo_url = format!("{}/tokeninfo", server.uri());
261 let _guard = EnvGuard::new("NBLM_TOKENINFO_ENDPOINT");
262 std::env::set_var("NBLM_TOKENINFO_ENDPOINT", &tokeninfo_url);
263
264 Mock::given(method("GET"))
265 .and(path("/tokeninfo"))
266 .and(query_param("access_token", "test-token"))
267 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
268 "scope": "https://www.googleapis.com/auth/cloud-platform"
269 })))
270 .expect(1)
271 .mount(&server)
272 .await;
273
274 Mock::given(method("POST"))
275 .and(path(
276 "/v1alpha/projects/123/locations/global/notebooks/notebook-id/sources:batchCreate",
277 ))
278 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
279 "sources": [],
280 "errorCount": 0
281 })))
282 .expect(0)
283 .mount(&server)
284 .await;
285
286 let client = build_client(&format!("{}/v1alpha", server.uri())).await;
287
288 let err = client
289 .add_sources(
290 "notebook-id",
291 vec![UserContent::google_drive(
292 "doc".to_string(),
293 "application/pdf".to_string(),
294 None,
295 )],
296 )
297 .await
298 .expect_err("expected add_sources to fail when drive scope is missing");
299
300 match err {
301 Error::TokenProvider(message) => {
302 assert!(
303 message.contains("drive.file"),
304 "unexpected message: {message}"
305 );
306 }
307 other => panic!("expected TokenProvider error, got {other:?}"),
308 }
309 }
310
311 #[tokio::test]
312 #[serial]
313 async fn batch_create_sources_validates_drive_scope() {
314 let server = MockServer::start().await;
315 let tokeninfo_url = format!("{}/tokeninfo", server.uri());
316 let _guard = EnvGuard::new("NBLM_TOKENINFO_ENDPOINT");
317 std::env::set_var("NBLM_TOKENINFO_ENDPOINT", &tokeninfo_url);
318
319 Mock::given(method("GET"))
320 .and(path("/tokeninfo"))
321 .and(query_param("access_token", "test-token"))
322 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
323 "scope": "https://www.googleapis.com/auth/drive"
324 })))
325 .expect(1)
326 .mount(&server)
327 .await;
328
329 Mock::given(method("POST"))
330 .and(path(
331 "/v1alpha/projects/123/locations/global/notebooks/notebook-id/sources:batchCreate",
332 ))
333 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
334 "sources": [],
335 "errorCount": 0
336 })))
337 .expect(1)
338 .mount(&server)
339 .await;
340
341 let client = build_client(&format!("{}/v1alpha", server.uri())).await;
342
343 let request = BatchCreateSourcesRequest {
344 user_contents: vec![UserContent::google_drive(
345 "doc".to_string(),
346 "application/pdf".to_string(),
347 None,
348 )],
349 };
350
351 let result = client
352 .batch_create_sources("notebook-id", request)
353 .await
354 .expect("expected batch_create_sources to succeed");
355
356 assert!(result.sources.is_empty());
357 }
358}