omni_dev/drive/
client.rs

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

omni_dev/drive/
client.rs
