omni_dev/drive/
client.rs

1//! Drive REST API client.
2//!
3//! A typed wrapper around [`GoogleApiClient`], the shared Google transport
4//! (`crate::drive::api_client`), pinned to the Drive v3 host. The transport
5//! attaches a Bearer access token refreshed by a shared
6//! [`DriveSession`](crate::drive::auth::DriveSession),
7//! retries HTTP 429 and Google's 403-shaped quota signal via the shared
8//! [`retry_if`](crate::utils::http::retry_if) driver, and retries exactly once
9//! on HTTP 401 by forcing a session refresh.
10//!
11//! Everything host-agnostic moved into `api_client.rs` in issue #1589 so the
12//! Sheets API — a *different host*, which a `/drive/v3/...` base URL cannot
13//! reach — could reuse it. What stays here is what is genuinely Drive's: the
14//! default host, the `DRIVE_API_URL` override, and the type identity that
15//! keeps `FilesApi` from being pointed at the wrong API.
16
17use anyhow::Result;
18use reqwest::Response;
19use url::Url;
20
21use crate::drive::api_client::GoogleApiClient;
22use crate::drive::auth::DriveCredentials;
23use crate::drive::error::DriveError;
24use crate::utils::env::{EnvSource, SystemEnv};
25
26/// The `service` tag Drive's HTTP records carry in the request log.
27///
28/// `SheetsClient` deliberately uses this **same** tag: the mutation records
29/// are already hardcoded `service: "drive"`
30/// (`crate::request_log::build_drive_mutation_record`), so splitting the HTTP
31/// records off under a separate service would make `omni-dev log --service
32/// drive` stop covering half of one feature's traffic. The host stays visible
33/// in each record's URL.
34pub(crate) const SERVICE_TAG: &str = "drive";
35
36/// Human-readable API name for Drive error messages. Unlike
37/// [`SERVICE_TAG`], which both clients share, this differs per API so a
38/// Sheets failure doesn't announce itself as a Drive one.
39pub(crate) const API_NAME: &str = "Drive";
40
41/// HTTP client for the Drive v3 REST API.
42pub struct DriveClient {
43    inner: GoogleApiClient,
44}
45
46impl std::fmt::Debug for DriveClient {
47    // Hand-written, not derived: omits `session` entirely rather than
48    // relying on every nested `Secret` staying wrapped — the safest
49    // possible redaction is "not mentioned at all." Deriving would also
50    // print the inner client's field names, including `session`.
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("DriveClient")
53            .field("base_url", &self.inner.base_url())
54            .finish_non_exhaustive()
55    }
56}
57
58impl DriveClient {
59    /// The real Drive API host. Unlike Gmail, which lives on its own
60    /// dedicated `gmail.googleapis.com` subdomain, Drive API v3 lives under
61    /// the general `googleapis.com` host with a `/drive/v3/...` path prefix
62    /// on each endpoint (see [ADR-0069](../../docs/adrs/adr-0069.md) §6).
63    /// Overridable wholesale via `DRIVE_API_URL`
64    /// (`crate::drive::auth::DRIVE_API_URL`; see
65    /// [`Self::from_credentials_with`]). [`Self::new`]'s `base_url`
66    /// parameter is the lower-level seam both the override and tests go
67    /// through.
68    const DEFAULT_BASE_URL: &'static str = "https://www.googleapis.com";
69
70    /// Builds a client against `base_url` with already-loaded credentials.
71    ///
72    /// For production use, construct via [`Self::from_credentials`]; tests
73    /// pass a wiremock URL directly.
74    pub fn new(base_url: &str, credentials: &DriveCredentials) -> Result<Self> {
75        Ok(Self {
76            inner: GoogleApiClient::new(base_url, credentials, SERVICE_TAG, API_NAME)?,
77        })
78    }
79
80    /// Creates a client from stored credentials against the real Drive API
81    /// host.
82    ///
83    /// Respects `DRIVE_API_URL` as an optional override: when set (and
84    /// non-empty) in the process environment it replaces
85    /// [`Self::DEFAULT_BASE_URL`] wholesale — mirrors Gmail's `GMAIL_API_URL`
86    /// (PR #1466), used to exercise CLI output shapes without a real Google
87    /// Cloud project or to route through a forced egress proxy.
88    pub fn from_credentials(credentials: &DriveCredentials) -> Result<Self> {
89        Self::from_credentials_with(&SystemEnv, credentials)
90    }
91
92    /// [`from_credentials`](Self::from_credentials) over an injected
93    /// [`EnvSource`], so tests can exercise the `DRIVE_API_URL` override via
94    /// `MapEnv` without mutating the process environment.
95    pub(crate) fn from_credentials_with(
96        env: &impl EnvSource,
97        credentials: &DriveCredentials,
98    ) -> Result<Self> {
99        let base_url = env
100            .var(crate::drive::auth::DRIVE_API_URL)
101            .filter(|s| !s.is_empty())
102            .unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
103        Self::new(&base_url, credentials)
104    }
105
106    /// Returns the API base URL (without trailing slash).
107    #[must_use]
108    pub fn base_url(&self) -> &str {
109        self.inner.base_url()
110    }
111
112    /// The shared transport, for deriving a client against another Google
113    /// host that reuses this one's OAuth session and connection pool.
114    ///
115    /// `dead_code`-allowed until `SheetsClient` lands in the following commit
116    /// (see [`GoogleApiClient::derived`]).
117    #[allow(dead_code)]
118    pub(in crate::drive) fn transport(&self) -> &GoogleApiClient {
119        &self.inner
120    }
121
122    /// Builds an absolute API URL by joining `path` onto `base_url`.
123    ///
124    /// Takes `base_url` (rather than `&self`) so the free `build_*_url`
125    /// functions in the API façade modules — and their unit tests, which
126    /// pass literal base URLs — can call it without an instance.
127    pub(crate) fn api_url(base_url: &str, path: &str) -> Result<Url> {
128        GoogleApiClient::api_url(base_url, path)
129    }
130
131    /// Checks `response` for success and deserialises its JSON body into `T`.
132    pub(crate) async fn parse_response<T: serde::de::DeserializeOwned>(
133        &self,
134        response: Response,
135        context: &'static str,
136    ) -> Result<T> {
137        self.inner.parse_response(response, context).await
138    }
139
140    /// Sends an authenticated GET and deserialises the JSON body into `T`.
141    pub(crate) async fn get_parsed<T: serde::de::DeserializeOwned>(
142        &self,
143        url: &str,
144        context: &'static str,
145    ) -> Result<T> {
146        self.inner.get_parsed(url, context).await
147    }
148
149    /// Sends an authenticated GET request and returns the raw response.
150    ///
151    /// Retries exactly once on HTTP 401 by forcing a session refresh.
152    pub async fn get_json(&self, url: &str) -> Result<Response> {
153        self.inner.get_json(url).await
154    }
155
156    /// Sends an authenticated GET request without forcing an `Accept:
157    /// application/json` header, for endpoints that return raw bytes rather
158    /// than JSON (`files.export`, `files.get?alt=media`) — see
159    /// [`Self::get_json`] for the JSON counterpart. Retries exactly once on
160    /// HTTP 401, identically to `get_json`.
161    pub async fn get_bytes(&self, url: &str) -> Result<Response> {
162        self.inner.get_bytes(url).await
163    }
164
165    /// Sends an authenticated POST request with a JSON body and returns the
166    /// raw response.
167    pub async fn post_json<T: serde::Serialize + Sync + ?Sized>(
168        &self,
169        url: &str,
170        body: &T,
171    ) -> Result<Response> {
172        self.inner.post_json(url, body).await
173    }
174
175    /// Sends an authenticated PATCH request with a JSON body and returns the
176    /// raw response. Drive's `files.update` (rename/move) is the only PATCH
177    /// endpoint this client calls with a JSON body.
178    pub async fn patch_json<T: serde::Serialize + Sync + ?Sized>(
179        &self,
180        url: &str,
181        body: &T,
182    ) -> Result<Response> {
183        self.inner.patch_json(url, body).await
184    }
185
186    /// Sends an authenticated POST request with a raw byte body and returns
187    /// the raw response — for Drive's multipart-upload endpoint, whose
188    /// `multipart/related` body [`crate::drive::files_api::FilesApi::upload`]
189    /// hand-assembles (Drive's upload endpoint rejects the
190    /// `multipart/form-data` `reqwest::multipart::Form` would produce).
191    pub async fn post_bytes(&self, url: &str, body: &[u8], content_type: &str) -> Result<Response> {
192        self.inner.post_bytes(url, body, content_type).await
193    }
194
195    /// Sends an authenticated PATCH request with a raw byte body and
196    /// returns the raw response —
197    /// [`crate::drive::files_api::FilesApi::edit_content`]'s simple
198    /// media-only content replacement (`uploadType=media`, no multipart
199    /// envelope needed since there's no accompanying metadata change).
200    pub async fn patch_bytes(
201        &self,
202        url: &str,
203        body: &[u8],
204        content_type: &str,
205    ) -> Result<Response> {
206        self.inner.patch_bytes(url, body, content_type).await
207    }
208
209    /// Consumes a non-success response into a [`DriveError`].
210    ///
211    /// Parses Google's `{"error":{"message":...}}` envelope — in either the
212    /// legacy `errors[].reason` or the newer `status` spelling, see
213    /// `api_client::error_reason` — into a human message when present
214    /// (falls back to the raw body otherwise). Drive signals quota
215    /// exhaustion as **403** `userRateLimitExceeded`, not `429`; that shape
216    /// drives a retry, so this only sees the error once retries are
217    /// exhausted (or the reason didn't match).
218    pub async fn response_to_error(response: Response) -> DriveError {
219        GoogleApiClient::response_to_error(API_NAME, response).await
220    }
221}
222
223/// Test-only seam letting sibling API-façade test modules (which can't
224/// reach `DriveClient`'s private fields directly, unlike this module's own
225/// `tests` submodule) bootstrap a deterministic access token via wiremock.
226#[cfg(test)]
227#[allow(clippy::expect_used)]
228pub(crate) mod test_support {
229    use super::{DriveClient, GoogleApiClient, API_NAME, SERVICE_TAG};
230    use crate::drive::auth::DriveCredentials;
231
232    /// Replaces `client`'s transport with one whose session points at an
233    /// explicit token endpoint.
234    ///
235    /// Rebuilds rather than mutating a session in place: the session lives
236    /// behind an `Arc` that a derived client may share, and assigning a new
237    /// one post-hoc would leave any already-derived client pointed at the
238    /// real `oauth2.googleapis.com`. Callers that need a *pair* of clients
239    /// should build the Drive client through this first and derive after.
240    pub(crate) fn replace_session(
241        client: &mut DriveClient,
242        credentials: &DriveCredentials,
243        token_endpoint: &str,
244    ) {
245        let base_url = client.base_url().to_string();
246        client.inner = GoogleApiClient::new_with_token_endpoint(
247            &base_url,
248            credentials,
249            SERVICE_TAG,
250            API_NAME,
251            token_endpoint,
252        )
253        .expect("failed to build a test transport");
254    }
255}
256
257#[cfg(test)]
258#[allow(clippy::unwrap_used, clippy::expect_used)]
259mod tests {
260    use super::*;
261    use crate::drive::auth::DriveGrantedScopes;
262    use crate::utils::secret::Secret;
263
264    fn test_credentials() -> DriveCredentials {
265        DriveCredentials {
266            client_id: "client-1".to_string(),
267            client_secret: Secret::new("secret-1"),
268            refresh_token: Secret::new("refresh-1"),
269            scope: DriveGrantedScopes::READONLY,
270        }
271    }
272
273    #[test]
274    fn new_client_strips_trailing_slash() {
275        let client = DriveClient::new("https://www.googleapis.com/", &test_credentials()).unwrap();
276        assert_eq!(client.base_url(), "https://www.googleapis.com");
277    }
278
279    #[test]
280    fn new_client_preserves_clean_url() {
281        let client = DriveClient::new("https://www.googleapis.com", &test_credentials()).unwrap();
282        assert_eq!(client.base_url(), "https://www.googleapis.com");
283    }
284
285    #[test]
286    fn from_credentials_uses_drive_api_host() {
287        // Via a fresh MapEnv, not from_credentials()'s real SystemEnv — a
288        // stray DRIVE_API_URL in the actual process environment must not
289        // make this test flaky (mirrors the Datadog/Gmail precedent).
290        let env = crate::test_support::env::MapEnv::new();
291        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
292        assert_eq!(client.base_url(), "https://www.googleapis.com");
293    }
294
295    #[test]
296    fn from_credentials_honours_api_url_override() {
297        let env = crate::test_support::env::MapEnv::new().with(
298            crate::drive::auth::DRIVE_API_URL,
299            "http://proxy.example:8080",
300        );
301        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
302        assert_eq!(client.base_url(), "http://proxy.example:8080");
303    }
304
305    #[test]
306    fn from_credentials_ignores_empty_api_url_override() {
307        let env =
308            crate::test_support::env::MapEnv::new().with(crate::drive::auth::DRIVE_API_URL, "");
309        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
310        assert_eq!(client.base_url(), "https://www.googleapis.com");
311    }
312
313    #[test]
314    fn client_debug_never_mentions_session_field() {
315        let client = DriveClient::new("https://www.googleapis.com", &test_credentials()).unwrap();
316        let debug = format!("{client:?}");
317        assert!(!debug.contains("secret-1"));
318        assert!(!debug.contains("refresh-1"));
319        assert!(!debug.contains("session"));
320        assert!(debug.contains("DriveClient"));
321    }
322
323    /// Mounts a bootstrap token-endpoint mock at the same base URL as the
324    /// Drive API mock — `DriveSession` doesn't distinguish the two hosts in
325    /// these tests, so pointing the token endpoint at the wiremock server
326    /// too keeps the setup to one server per test.
327    async fn client_with_bootstrapped_token(server: &wiremock::MockServer) -> DriveClient {
328        // `up_to_n_times(1)` + `with_priority(1)` so a test's own follow-up
329        // POST /token mock (registered at `with_priority(2)`, matched only
330        // once this one is exhausted) can simulate a second, distinct
331        // refresh without either mock racing the other for every request.
332        wiremock::Mock::given(wiremock::matchers::method("POST"))
333            .and(wiremock::matchers::path("/token"))
334            .respond_with(
335                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
336                    "access_token": "bootstrap-token",
337                    "expires_in": 3600,
338                })),
339            )
340            .up_to_n_times(1)
341            .with_priority(1)
342            .mount(server)
343            .await;
344
345        let mut client = DriveClient::new(&server.uri(), &test_credentials()).unwrap();
346        test_support::replace_session(
347            &mut client,
348            &test_credentials(),
349            &format!("{}/token", server.uri()),
350        );
351        client
352    }
353    #[tokio::test]
354    async fn get_json_sends_bearer_auth_header() {
355        let server = wiremock::MockServer::start().await;
356        let client = client_with_bootstrapped_token(&server).await;
357        wiremock::Mock::given(wiremock::matchers::method("GET"))
358            .and(wiremock::matchers::path("/test"))
359            .and(wiremock::matchers::header(
360                "Authorization",
361                "Bearer bootstrap-token",
362            ))
363            .respond_with(
364                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true})),
365            )
366            .expect(1)
367            .mount(&server)
368            .await;
369
370        let resp = client
371            .get_json(&format!("{}/test", server.uri()))
372            .await
373            .unwrap();
374        assert!(resp.status().is_success());
375    }
376
377    #[tokio::test]
378    async fn get_bytes_sends_bearer_auth_header_without_json_accept_header() {
379        let server = wiremock::MockServer::start().await;
380        let client = client_with_bootstrapped_token(&server).await;
381        wiremock::Mock::given(wiremock::matchers::method("GET"))
382            .and(wiremock::matchers::path("/test"))
383            .and(wiremock::matchers::header(
384                "Authorization",
385                "Bearer bootstrap-token",
386            ))
387            .respond_with(
388                wiremock::ResponseTemplate::new(200).set_body_bytes(b"raw bytes".to_vec()),
389            )
390            .expect(1)
391            .mount(&server)
392            .await;
393
394        let resp = client
395            .get_bytes(&format!("{}/test", server.uri()))
396            .await
397            .unwrap();
398        assert!(resp.status().is_success());
399        let bytes = resp.bytes().await.unwrap();
400        assert_eq!(bytes.as_ref(), b"raw bytes");
401    }
402
403    #[tokio::test]
404    async fn get_bytes_retries_on_429() {
405        let server = wiremock::MockServer::start().await;
406        let client = client_with_bootstrapped_token(&server).await;
407        wiremock::Mock::given(wiremock::matchers::method("GET"))
408            .and(wiremock::matchers::path("/test"))
409            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
410            .up_to_n_times(1)
411            .with_priority(1)
412            .mount(&server)
413            .await;
414        wiremock::Mock::given(wiremock::matchers::method("GET"))
415            .and(wiremock::matchers::path("/test"))
416            .respond_with(wiremock::ResponseTemplate::new(200).set_body_bytes(b"ok".to_vec()))
417            .with_priority(2)
418            .mount(&server)
419            .await;
420
421        let resp = client
422            .get_bytes(&format!("{}/test", server.uri()))
423            .await
424            .unwrap();
425        assert_eq!(resp.status().as_u16(), 200);
426    }
427
428    #[tokio::test]
429    async fn get_bytes_refreshes_and_retries_once_on_401() {
430        let server = wiremock::MockServer::start().await;
431        let client = client_with_bootstrapped_token(&server).await;
432        wiremock::Mock::given(wiremock::matchers::method("POST"))
433            .and(wiremock::matchers::path("/token"))
434            .respond_with(
435                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
436                    "access_token": "refreshed-token",
437                    "expires_in": 3600,
438                })),
439            )
440            .up_to_n_times(1)
441            .with_priority(2)
442            .mount(&server)
443            .await;
444        wiremock::Mock::given(wiremock::matchers::method("GET"))
445            .and(wiremock::matchers::path("/test"))
446            .and(wiremock::matchers::header(
447                "Authorization",
448                "Bearer bootstrap-token",
449            ))
450            .respond_with(wiremock::ResponseTemplate::new(401))
451            .expect(1)
452            .mount(&server)
453            .await;
454        wiremock::Mock::given(wiremock::matchers::method("GET"))
455            .and(wiremock::matchers::path("/test"))
456            .and(wiremock::matchers::header(
457                "Authorization",
458                "Bearer refreshed-token",
459            ))
460            .respond_with(wiremock::ResponseTemplate::new(200).set_body_bytes(b"ok".to_vec()))
461            .expect(1)
462            .mount(&server)
463            .await;
464
465        let resp = client
466            .get_bytes(&format!("{}/test", server.uri()))
467            .await
468            .unwrap();
469        assert_eq!(resp.status().as_u16(), 200);
470    }
471
472    #[tokio::test]
473    async fn get_bytes_propagates_network_errors() {
474        let client = DriveClient::new("http://127.0.0.1:1", &test_credentials()).unwrap();
475        let result = client.get_bytes("http://127.0.0.1:1/test").await;
476        assert!(result.is_err());
477    }
478
479    #[tokio::test]
480    async fn post_json_sends_body_and_bearer_auth() {
481        let server = wiremock::MockServer::start().await;
482        let client = client_with_bootstrapped_token(&server).await;
483        wiremock::Mock::given(wiremock::matchers::method("POST"))
484            .and(wiremock::matchers::path("/test"))
485            .and(wiremock::matchers::header(
486                "Authorization",
487                "Bearer bootstrap-token",
488            ))
489            .and(wiremock::matchers::body_json(serde_json::json!({"k": "v"})))
490            .respond_with(wiremock::ResponseTemplate::new(200))
491            .expect(1)
492            .mount(&server)
493            .await;
494
495        let resp = client
496            .post_json(
497                &format!("{}/test", server.uri()),
498                &serde_json::json!({"k": "v"}),
499            )
500            .await
501            .unwrap();
502        assert!(resp.status().is_success());
503    }
504
505    #[tokio::test]
506    async fn patch_json_sends_body_and_bearer_auth() {
507        let server = wiremock::MockServer::start().await;
508        let client = client_with_bootstrapped_token(&server).await;
509        wiremock::Mock::given(wiremock::matchers::method("PATCH"))
510            .and(wiremock::matchers::path("/test"))
511            .and(wiremock::matchers::header(
512                "Authorization",
513                "Bearer bootstrap-token",
514            ))
515            .and(wiremock::matchers::body_json(
516                serde_json::json!({"name": "new-name"}),
517            ))
518            .respond_with(wiremock::ResponseTemplate::new(200))
519            .expect(1)
520            .mount(&server)
521            .await;
522
523        let resp = client
524            .patch_json(
525                &format!("{}/test", server.uri()),
526                &serde_json::json!({"name": "new-name"}),
527            )
528            .await
529            .unwrap();
530        assert!(resp.status().is_success());
531    }
532
533    #[tokio::test]
534    async fn get_json_retries_on_429() {
535        let server = wiremock::MockServer::start().await;
536        let client = client_with_bootstrapped_token(&server).await;
537        wiremock::Mock::given(wiremock::matchers::method("GET"))
538            .and(wiremock::matchers::path("/test"))
539            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
540            .up_to_n_times(1)
541            .with_priority(1)
542            .mount(&server)
543            .await;
544        wiremock::Mock::given(wiremock::matchers::method("GET"))
545            .and(wiremock::matchers::path("/test"))
546            .respond_with(wiremock::ResponseTemplate::new(200))
547            .with_priority(2)
548            .mount(&server)
549            .await;
550
551        let resp = client
552            .get_json(&format!("{}/test", server.uri()))
553            .await
554            .unwrap();
555        assert_eq!(resp.status().as_u16(), 200);
556    }
557
558    #[tokio::test]
559    async fn get_json_retries_403_user_rate_limit_exceeded_then_succeeds() {
560        let server = wiremock::MockServer::start().await;
561        let client = client_with_bootstrapped_token(&server).await;
562        wiremock::Mock::given(wiremock::matchers::method("GET"))
563            .and(wiremock::matchers::path("/test"))
564            .respond_with(
565                wiremock::ResponseTemplate::new(403)
566                    .append_header("Retry-After", "0")
567                    .set_body_json(serde_json::json!({
568                        "error": {"message": "User Rate Limit Exceeded", "errors": [{"reason": "userRateLimitExceeded"}]}
569                    })),
570            )
571            .up_to_n_times(1)
572            .with_priority(1)
573            .mount(&server)
574            .await;
575        wiremock::Mock::given(wiremock::matchers::method("GET"))
576            .and(wiremock::matchers::path("/test"))
577            .respond_with(wiremock::ResponseTemplate::new(200))
578            .with_priority(2)
579            .mount(&server)
580            .await;
581
582        let resp = client
583            .get_json(&format!("{}/test", server.uri()))
584            .await
585            .unwrap();
586        assert_eq!(resp.status().as_u16(), 200);
587    }
588
589    #[tokio::test]
590    async fn get_json_does_not_retry_insufficient_permissions_403() {
591        let server = wiremock::MockServer::start().await;
592        let client = client_with_bootstrapped_token(&server).await;
593        wiremock::Mock::given(wiremock::matchers::method("GET"))
594            .and(wiremock::matchers::path("/test"))
595            .respond_with(
596                wiremock::ResponseTemplate::new(403).set_body_json(serde_json::json!({
597                    "error": {"message": "Insufficient Permission", "errors": [{"reason": "insufficientPermissions"}]}
598                })),
599            )
600            .expect(1)
601            .mount(&server)
602            .await;
603
604        let resp = client
605            .get_json(&format!("{}/test", server.uri()))
606            .await
607            .unwrap();
608        assert_eq!(resp.status().as_u16(), 403);
609    }
610
611    #[tokio::test]
612    async fn get_json_refreshes_and_retries_once_on_401() {
613        let server = wiremock::MockServer::start().await;
614        let client = client_with_bootstrapped_token(&server).await;
615        // The refresh endpoint issues a second, distinct token.
616        wiremock::Mock::given(wiremock::matchers::method("POST"))
617            .and(wiremock::matchers::path("/token"))
618            .respond_with(
619                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
620                    "access_token": "refreshed-token",
621                    "expires_in": 3600,
622                })),
623            )
624            .up_to_n_times(1)
625            .with_priority(2)
626            .mount(&server)
627            .await;
628        wiremock::Mock::given(wiremock::matchers::method("GET"))
629            .and(wiremock::matchers::path("/test"))
630            .and(wiremock::matchers::header(
631                "Authorization",
632                "Bearer bootstrap-token",
633            ))
634            .respond_with(wiremock::ResponseTemplate::new(401))
635            .expect(1)
636            .mount(&server)
637            .await;
638        wiremock::Mock::given(wiremock::matchers::method("GET"))
639            .and(wiremock::matchers::path("/test"))
640            .and(wiremock::matchers::header(
641                "Authorization",
642                "Bearer refreshed-token",
643            ))
644            .respond_with(wiremock::ResponseTemplate::new(200))
645            .expect(1)
646            .mount(&server)
647            .await;
648
649        let resp = client
650            .get_json(&format!("{}/test", server.uri()))
651            .await
652            .unwrap();
653        assert_eq!(resp.status().as_u16(), 200);
654    }
655
656    #[tokio::test]
657    async fn get_json_does_not_retry_a_second_time_on_persistent_401() {
658        let server = wiremock::MockServer::start().await;
659        let client = client_with_bootstrapped_token(&server).await;
660        wiremock::Mock::given(wiremock::matchers::method("POST"))
661            .and(wiremock::matchers::path("/token"))
662            .respond_with(
663                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
664                    "access_token": "still-rejected-token",
665                    "expires_in": 3600,
666                })),
667            )
668            .up_to_n_times(1)
669            .with_priority(2)
670            .mount(&server)
671            .await;
672        wiremock::Mock::given(wiremock::matchers::method("GET"))
673            .and(wiremock::matchers::path("/test"))
674            .respond_with(
675                wiremock::ResponseTemplate::new(401).set_body_string("still unauthorized"),
676            )
677            .expect(2)
678            .mount(&server)
679            .await;
680
681        let resp = client
682            .get_json(&format!("{}/test", server.uri()))
683            .await
684            .unwrap();
685        assert_eq!(resp.status().as_u16(), 401);
686    }
687
688    #[tokio::test]
689    async fn response_to_error_extracts_drive_message_and_reason() {
690        let server = wiremock::MockServer::start().await;
691        let client = client_with_bootstrapped_token(&server).await;
692        // `userRateLimitExceeded` is now retryable (`is_drive_quota_exceeded`),
693        // so without a zero-delay `Retry-After` this test would wait through
694        // the real exponential backoff before giving up.
695        wiremock::Mock::given(wiremock::matchers::method("GET"))
696            .and(wiremock::matchers::path("/test"))
697            .respond_with(
698                wiremock::ResponseTemplate::new(403)
699                    .append_header("Retry-After", "0")
700                    .set_body_json(serde_json::json!({
701                        "error": {
702                            "message": "User Rate Limit Exceeded",
703                            "errors": [{"reason": "userRateLimitExceeded"}],
704                        }
705                    })),
706            )
707            .mount(&server)
708            .await;
709
710        let resp = client
711            .get_json(&format!("{}/test", server.uri()))
712            .await
713            .unwrap();
714        let err = DriveClient::response_to_error(resp).await;
715        let msg = err.to_string();
716        assert!(msg.contains("User Rate Limit Exceeded"));
717        assert!(msg.contains("userRateLimitExceeded"));
718        assert_eq!(err.reason(), Some("userRateLimitExceeded"));
719    }
720
721    #[tokio::test]
722    async fn response_to_error_omits_reason_suffix_when_absent() {
723        let server = wiremock::MockServer::start().await;
724        let client = client_with_bootstrapped_token(&server).await;
725        wiremock::Mock::given(wiremock::matchers::method("GET"))
726            .and(wiremock::matchers::path("/test"))
727            .respond_with(
728                wiremock::ResponseTemplate::new(400).set_body_json(serde_json::json!({
729                    "error": {
730                        "message": "Invalid request",
731                    }
732                })),
733            )
734            .mount(&server)
735            .await;
736
737        let resp = client
738            .get_json(&format!("{}/test", server.uri()))
739            .await
740            .unwrap();
741        let err = DriveClient::response_to_error(resp).await;
742        let msg = err.to_string();
743        assert!(msg.contains("Invalid request"));
744        assert!(!msg.contains("reason:"));
745        assert_eq!(err.reason(), None);
746    }
747
748    #[tokio::test]
749    async fn response_to_error_falls_back_to_raw_body_when_not_drive_shaped() {
750        let server = wiremock::MockServer::start().await;
751        let client = client_with_bootstrapped_token(&server).await;
752        wiremock::Mock::given(wiremock::matchers::method("GET"))
753            .and(wiremock::matchers::path("/test"))
754            .respond_with(wiremock::ResponseTemplate::new(500).set_body_string("internal error"))
755            .mount(&server)
756            .await;
757
758        let resp = client
759            .get_json(&format!("{}/test", server.uri()))
760            .await
761            .unwrap();
762        let err = DriveClient::response_to_error(resp).await;
763        assert!(err.to_string().contains("internal error"));
764    }
765
766    #[tokio::test]
767    async fn get_json_propagates_network_errors() {
768        let client = DriveClient::new("http://127.0.0.1:1", &test_credentials()).unwrap();
769        let result = client.get_json("http://127.0.0.1:1/test").await;
770        assert!(result.is_err());
771    }
772
773    #[tokio::test]
774    async fn get_parsed_errors_on_malformed_json_response() {
775        let server = wiremock::MockServer::start().await;
776        let client = client_with_bootstrapped_token(&server).await;
777        wiremock::Mock::given(wiremock::matchers::method("GET"))
778            .and(wiremock::matchers::path("/test"))
779            .respond_with(wiremock::ResponseTemplate::new(200).set_body_string("not json"))
780            .mount(&server)
781            .await;
782
783        let result: Result<serde_json::Value> = client
784            .get_parsed(&format!("{}/test", server.uri()), "test context")
785            .await;
786        assert!(result.is_err());
787    }
788
789    #[tokio::test]
790    async fn get_parsed_errors_on_non_success_status_without_parsing_the_body() {
791        let server = wiremock::MockServer::start().await;
792        let client = client_with_bootstrapped_token(&server).await;
793        wiremock::Mock::given(wiremock::matchers::method("GET"))
794            .and(wiremock::matchers::path("/test"))
795            .respond_with(
796                wiremock::ResponseTemplate::new(404).set_body_json(serde_json::json!({
797                    "error": {"message": "File not found"}
798                })),
799            )
800            .mount(&server)
801            .await;
802
803        let result: Result<serde_json::Value> = client
804            .get_parsed(&format!("{}/test", server.uri()), "test context")
805            .await;
806        let err = result.unwrap_err();
807        assert!(err.to_string().contains("File not found"));
808    }
809
810    #[tokio::test]
811    async fn put_json_sends_body_and_bearer_auth() {
812        let server = wiremock::MockServer::start().await;
813        let client = client_with_bootstrapped_token(&server).await;
814        wiremock::Mock::given(wiremock::matchers::method("PUT"))
815            .and(wiremock::matchers::path("/test"))
816            .and(wiremock::matchers::header(
817                "Authorization",
818                "Bearer bootstrap-token",
819            ))
820            .and(wiremock::matchers::body_json(serde_json::json!({
821                "values": [["a", "b"]],
822            })))
823            .respond_with(wiremock::ResponseTemplate::new(200).set_body_string("{}"))
824            .expect(1)
825            .mount(&server)
826            .await;
827
828        let response = client
829            .transport()
830            .put_json(
831                &format!("{}/test", server.uri()),
832                &serde_json::json!({"values": [["a", "b"]]}),
833            )
834            .await
835            .unwrap();
836        assert_eq!(response.status().as_u16(), 200);
837    }
838
839    #[tokio::test]
840    async fn a_drive_failure_still_names_drive() {
841        // The other half of the pair: threading an API name must not have
842        // changed what Drive's own failures say.
843        let server = wiremock::MockServer::start().await;
844        let client = client_with_bootstrapped_token(&server).await;
845        wiremock::Mock::given(wiremock::matchers::method("GET"))
846            .and(wiremock::matchers::path("/test"))
847            .respond_with(wiremock::ResponseTemplate::new(500).set_body_string("boom"))
848            .mount(&server)
849            .await;
850
851        let response = client
852            .get_json(&format!("{}/test", server.uri()))
853            .await
854            .unwrap();
855        let text = DriveClient::response_to_error(response).await.to_string();
856        assert!(text.starts_with("Drive API request failed"), "{text}");
857    }
858
859    #[tokio::test]
860    async fn response_to_error_reads_the_google_rpc_envelope_sheets_returns() {
861        let server = wiremock::MockServer::start().await;
862        let client = client_with_bootstrapped_token(&server).await;
863        wiremock::Mock::given(wiremock::matchers::method("GET"))
864            .and(wiremock::matchers::path("/test"))
865            .respond_with(
866                wiremock::ResponseTemplate::new(403).set_body_json(serde_json::json!({
867                    "error": {
868                        "code": 403,
869                        "message": "The caller does not have permission",
870                        "status": "PERMISSION_DENIED",
871                    },
872                })),
873            )
874            .mount(&server)
875            .await;
876
877        let response = client
878            .get_json(&format!("{}/test", server.uri()))
879            .await
880            .unwrap();
881        let error = DriveClient::response_to_error(response).await;
882        let DriveError::ApiRequestFailed { reason, .. } = error else {
883            panic!("expected ApiRequestFailed");
884        };
885        assert_eq!(reason.as_deref(), Some("PERMISSION_DENIED"));
886    }
887}
omni_dev/drive/client.rs

omni_dev/drive/
client.rs
