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 with a JSON body.
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 an authenticated POST request with a raw byte body and returns
195    /// the raw response — for Drive's multipart-upload endpoint, whose
196    /// `multipart/related` body [`crate::drive::files_api::FilesApi::upload`]
197    /// hand-assembles (Drive's upload endpoint rejects the
198    /// `multipart/form-data` `reqwest::multipart::Form` would produce).
199    ///
200    /// `body` is cloned per send attempt (`send_authorized`'s `build`
201    /// closure is `Fn`, not `FnOnce` — it may run twice, once on a 401
202    /// retry — and [`reqwest::RequestBuilder::body`] takes ownership).
203    pub async fn post_bytes(&self, url: &str, body: &[u8], content_type: &str) -> Result<Response> {
204        self.send_authorized(url, "POST", |client, token| {
205            client
206                .post(url)
207                .bearer_auth(token)
208                .header("Content-Type", content_type)
209                .body(body.to_vec())
210        })
211        .await
212    }
213
214    /// Sends an authenticated PATCH request with a raw byte body and
215    /// returns the raw response —
216    /// [`crate::drive::files_api::FilesApi::edit_content`]'s simple
217    /// media-only content replacement (`uploadType=media`, no multipart
218    /// envelope needed since there's no accompanying metadata change).
219    pub async fn patch_bytes(
220        &self,
221        url: &str,
222        body: &[u8],
223        content_type: &str,
224    ) -> Result<Response> {
225        self.send_authorized(url, "PATCH", |client, token| {
226            client
227                .patch(url)
228                .bearer_auth(token)
229                .header("Content-Type", content_type)
230                .body(body.to_vec())
231        })
232        .await
233    }
234
235    /// Sends a request built by `build`, retrying exactly once on HTTP 401.
236    ///
237    /// [`DriveSession::access_token`] already refreshes proactively when the
238    /// tracked expiry is near — this reactive path exists for what
239    /// proactive tracking can't see: clock skew against Google's clock, or
240    /// the token being invalidated server-side mid-run (revoked access). A
241    /// second 401 after the retry is authoritative: either the refresh
242    /// produced a token that was also rejected, or another caller's
243    /// already-current token was reused and still rejected — either way the
244    /// problem isn't staleness, so it surfaces as an ordinary
245    /// `ApiRequestFailed` rather than retrying again.
246    async fn send_authorized<F>(
247        &self,
248        url: &str,
249        method: &'static str,
250        build: F,
251    ) -> Result<Response>
252    where
253        F: Fn(&Client, &str) -> reqwest::RequestBuilder + Send + Sync,
254    {
255        let token = self
256            .session
257            .access_token()
258            .await
259            .context("Failed to obtain a Drive access token")?;
260        let response = self
261            .send_once(url, method, &build, token.expose_secret())
262            .await?;
263        if response.status().as_u16() != 401 {
264            return Ok(response);
265        }
266        let refreshed = self
267            .session
268            .force_refresh(&token)
269            .await
270            .context("Failed to refresh the Drive access token after a 401")?;
271        self.send_once(url, method, &build, refreshed.expose_secret())
272            .await
273    }
274
275    async fn send_once<F>(
276        &self,
277        url: &str,
278        method: &'static str,
279        build: &F,
280        token: &str,
281    ) -> Result<Response>
282    where
283        F: Fn(&Client, &str) -> reqwest::RequestBuilder + Send + Sync,
284    {
285        retry_if(
286            || build(&self.client, token),
287            |started, result| {
288                request_log::record_http_result("drive", method, url, started, result);
289            },
290            |status, body| status == 429 || is_drive_quota_exceeded(status, body),
291        )
292        .await
293        .with_context(|| format!("Failed to send {method} request to Drive API"))
294    }
295
296    /// Consumes a non-success response into a [`DriveError`].
297    ///
298    /// Parses Drive's `{"error":{"message":...,"errors":[{"reason":...}]}}`
299    /// envelope (the same shape every Google API error uses) into a human
300    /// message when present (falls back to the raw body otherwise). Drive
301    /// signals quota exhaustion as **403** `userRateLimitExceeded`, not
302    /// `429` — unlike plain 429s, that shape now also drives a retry (see
303    /// [`is_drive_quota_exceeded`] via [`retry_if`](crate::utils::http::retry_if)),
304    /// so this only sees the error once retries are exhausted (or the
305    /// reason didn't match).
306    pub async fn response_to_error(response: Response) -> DriveError {
307        let status = response.status().as_u16();
308        let raw = response.text().await.unwrap_or_default();
309        let value = serde_json::from_str::<serde_json::Value>(&raw).ok();
310        let reason = value.as_ref().and_then(drive_error_reason);
311        let body = value
312            .as_ref()
313            .and_then(drive_error_message)
314            .map(|message| match &reason {
315                Some(r) => format!("{message} (reason: {r})"),
316                None => message,
317            })
318            .unwrap_or(raw);
319        DriveError::ApiRequestFailed {
320            status,
321            body,
322            reason,
323        }
324    }
325}
326
327/// Extracts the `error.errors[0].reason` field from Drive's already-parsed
328/// JSON error envelope, if present.
329fn drive_error_reason(value: &serde_json::Value) -> Option<String> {
330    value
331        .get("error")
332        .and_then(|e| e.get("errors"))
333        .and_then(|e| e.as_array())
334        .and_then(|a| a.first())
335        .and_then(|e| e.get("reason"))
336        .and_then(|r| r.as_str())
337        .map(str::to_string)
338}
339
340/// Extracts the `error.message` field from Drive's already-parsed JSON
341/// error envelope, if present.
342fn drive_error_message(value: &serde_json::Value) -> Option<String> {
343    value
344        .get("error")?
345        .get("message")?
346        .as_str()
347        .map(str::to_string)
348}
349
350/// Whether a response is Drive's quota-exhaustion signal — **403** with
351/// `reason` of `userRateLimitExceeded` specifically, not any 403 with a
352/// `reason` field: e.g. `insufficientPermissions` is also a 403 and must
353/// never be retried (retrying a scope/permission error just wastes the
354/// backoff window before failing anyway).
355///
356/// Unlike Gmail's `is_gmail_quota_exceeded`, this does **not** also match
357/// the bare `rateLimitExceeded` reason — that string is confirmed for
358/// Gmail but not confirmed for Drive against
359/// [Drive's error-handling guide](https://developers.google.com/workspace/drive/api/guides/handle-errors);
360/// widen this match only once testing surfaces it. Plain `429` responses
361/// (Drive's other documented rate-limit signal) are already covered by the
362/// literal `status == 429` branch in [`DriveClient::send_once`], independent
363/// of this function.
364fn is_drive_quota_exceeded(status: u16, body: &[u8]) -> bool {
365    if status != 403 {
366        return false;
367    }
368    let Ok(text) = std::str::from_utf8(body) else {
369        return false;
370    };
371    let Ok(value) = serde_json::from_str::<serde_json::Value>(text) else {
372        return false;
373    };
374    matches!(
375        drive_error_reason(&value).as_deref(),
376        Some("userRateLimitExceeded")
377    )
378}
379
380/// Test-only seam letting sibling API-façade test modules (which can't
381/// reach `DriveClient`'s private fields directly, unlike this module's own
382/// `tests` submodule) bootstrap a deterministic access token via wiremock.
383#[cfg(test)]
384pub(crate) mod test_support {
385    use super::DriveClient;
386    use crate::drive::auth::{DriveCredentials, DriveSession};
387
388    /// Replaces `client`'s session with one pointed at an explicit token
389    /// endpoint.
390    pub(crate) fn replace_session(
391        client: &mut DriveClient,
392        credentials: &DriveCredentials,
393        token_endpoint: &str,
394    ) {
395        client.session = DriveSession::new_with_token_endpoint(
396            client.client.clone(),
397            credentials,
398            token_endpoint,
399        );
400    }
401}
402
403#[cfg(test)]
404#[allow(clippy::unwrap_used, clippy::expect_used)]
405mod tests {
406    use super::*;
407    use crate::drive::auth::DriveGrantedScopes;
408    use crate::utils::secret::Secret;
409
410    fn test_credentials() -> DriveCredentials {
411        DriveCredentials {
412            client_id: "client-1".to_string(),
413            client_secret: Secret::new("secret-1"),
414            refresh_token: Secret::new("refresh-1"),
415            scope: DriveGrantedScopes::READONLY,
416        }
417    }
418
419    #[test]
420    fn is_drive_quota_exceeded_false_on_non_utf8_body() {
421        assert!(!is_drive_quota_exceeded(403, &[0xff, 0xfe]));
422    }
423
424    #[test]
425    fn is_drive_quota_exceeded_false_on_non_403_status() {
426        assert!(!is_drive_quota_exceeded(
427            429,
428            br#"{"error":{"errors":[{"reason":"userRateLimitExceeded"}]}}"#
429        ));
430    }
431
432    #[test]
433    fn is_drive_quota_exceeded_true_on_matching_403_reason() {
434        assert!(is_drive_quota_exceeded(
435            403,
436            br#"{"error":{"errors":[{"reason":"userRateLimitExceeded"}]}}"#
437        ));
438    }
439
440    #[test]
441    fn new_client_strips_trailing_slash() {
442        let client = DriveClient::new("https://www.googleapis.com/", &test_credentials()).unwrap();
443        assert_eq!(client.base_url(), "https://www.googleapis.com");
444    }
445
446    #[test]
447    fn new_client_preserves_clean_url() {
448        let client = DriveClient::new("https://www.googleapis.com", &test_credentials()).unwrap();
449        assert_eq!(client.base_url(), "https://www.googleapis.com");
450    }
451
452    #[test]
453    fn from_credentials_uses_drive_api_host() {
454        // Via a fresh MapEnv, not from_credentials()'s real SystemEnv — a
455        // stray DRIVE_API_URL in the actual process environment must not
456        // make this test flaky (mirrors the Datadog/Gmail precedent).
457        let env = crate::test_support::env::MapEnv::new();
458        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
459        assert_eq!(client.base_url(), "https://www.googleapis.com");
460    }
461
462    #[test]
463    fn from_credentials_honours_api_url_override() {
464        let env = crate::test_support::env::MapEnv::new().with(
465            crate::drive::auth::DRIVE_API_URL,
466            "http://proxy.example:8080",
467        );
468        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
469        assert_eq!(client.base_url(), "http://proxy.example:8080");
470    }
471
472    #[test]
473    fn from_credentials_ignores_empty_api_url_override() {
474        let env =
475            crate::test_support::env::MapEnv::new().with(crate::drive::auth::DRIVE_API_URL, "");
476        let client = DriveClient::from_credentials_with(&env, &test_credentials()).unwrap();
477        assert_eq!(client.base_url(), "https://www.googleapis.com");
478    }
479
480    #[test]
481    fn client_debug_never_mentions_session_field() {
482        let client = DriveClient::new("https://www.googleapis.com", &test_credentials()).unwrap();
483        let debug = format!("{client:?}");
484        assert!(!debug.contains("secret-1"));
485        assert!(!debug.contains("refresh-1"));
486        assert!(!debug.contains("session"));
487        assert!(debug.contains("DriveClient"));
488    }
489
490    /// Mounts a bootstrap token-endpoint mock at the same base URL as the
491    /// Drive API mock — `DriveSession` doesn't distinguish the two hosts in
492    /// these tests, so pointing the token endpoint at the wiremock server
493    /// too keeps the setup to one server per test.
494    async fn client_with_bootstrapped_token(server: &wiremock::MockServer) -> DriveClient {
495        // `up_to_n_times(1)` + `with_priority(1)` so a test's own follow-up
496        // POST /token mock (registered at `with_priority(2)`, matched only
497        // once this one is exhausted) can simulate a second, distinct
498        // refresh without either mock racing the other for every request.
499        wiremock::Mock::given(wiremock::matchers::method("POST"))
500            .and(wiremock::matchers::path("/token"))
501            .respond_with(
502                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
503                    "access_token": "bootstrap-token",
504                    "expires_in": 3600,
505                })),
506            )
507            .up_to_n_times(1)
508            .with_priority(1)
509            .mount(server)
510            .await;
511
512        let mut client = DriveClient::new(&server.uri(), &test_credentials()).unwrap();
513        client.session = DriveSession::new_with_token_endpoint(
514            client.client.clone(),
515            &test_credentials(),
516            &format!("{}/token", server.uri()),
517        );
518        client
519    }
520
521    #[tokio::test]
522    async fn get_json_sends_bearer_auth_header() {
523        let server = wiremock::MockServer::start().await;
524        let client = client_with_bootstrapped_token(&server).await;
525        wiremock::Mock::given(wiremock::matchers::method("GET"))
526            .and(wiremock::matchers::path("/test"))
527            .and(wiremock::matchers::header(
528                "Authorization",
529                "Bearer bootstrap-token",
530            ))
531            .respond_with(
532                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true})),
533            )
534            .expect(1)
535            .mount(&server)
536            .await;
537
538        let resp = client
539            .get_json(&format!("{}/test", server.uri()))
540            .await
541            .unwrap();
542        assert!(resp.status().is_success());
543    }
544
545    #[tokio::test]
546    async fn get_bytes_sends_bearer_auth_header_without_json_accept_header() {
547        let server = wiremock::MockServer::start().await;
548        let client = client_with_bootstrapped_token(&server).await;
549        wiremock::Mock::given(wiremock::matchers::method("GET"))
550            .and(wiremock::matchers::path("/test"))
551            .and(wiremock::matchers::header(
552                "Authorization",
553                "Bearer bootstrap-token",
554            ))
555            .respond_with(
556                wiremock::ResponseTemplate::new(200).set_body_bytes(b"raw bytes".to_vec()),
557            )
558            .expect(1)
559            .mount(&server)
560            .await;
561
562        let resp = client
563            .get_bytes(&format!("{}/test", server.uri()))
564            .await
565            .unwrap();
566        assert!(resp.status().is_success());
567        let bytes = resp.bytes().await.unwrap();
568        assert_eq!(bytes.as_ref(), b"raw bytes");
569    }
570
571    #[tokio::test]
572    async fn get_bytes_retries_on_429() {
573        let server = wiremock::MockServer::start().await;
574        let client = client_with_bootstrapped_token(&server).await;
575        wiremock::Mock::given(wiremock::matchers::method("GET"))
576            .and(wiremock::matchers::path("/test"))
577            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
578            .up_to_n_times(1)
579            .with_priority(1)
580            .mount(&server)
581            .await;
582        wiremock::Mock::given(wiremock::matchers::method("GET"))
583            .and(wiremock::matchers::path("/test"))
584            .respond_with(wiremock::ResponseTemplate::new(200).set_body_bytes(b"ok".to_vec()))
585            .with_priority(2)
586            .mount(&server)
587            .await;
588
589        let resp = client
590            .get_bytes(&format!("{}/test", server.uri()))
591            .await
592            .unwrap();
593        assert_eq!(resp.status().as_u16(), 200);
594    }
595
596    #[tokio::test]
597    async fn get_bytes_refreshes_and_retries_once_on_401() {
598        let server = wiremock::MockServer::start().await;
599        let client = client_with_bootstrapped_token(&server).await;
600        wiremock::Mock::given(wiremock::matchers::method("POST"))
601            .and(wiremock::matchers::path("/token"))
602            .respond_with(
603                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
604                    "access_token": "refreshed-token",
605                    "expires_in": 3600,
606                })),
607            )
608            .up_to_n_times(1)
609            .with_priority(2)
610            .mount(&server)
611            .await;
612        wiremock::Mock::given(wiremock::matchers::method("GET"))
613            .and(wiremock::matchers::path("/test"))
614            .and(wiremock::matchers::header(
615                "Authorization",
616                "Bearer bootstrap-token",
617            ))
618            .respond_with(wiremock::ResponseTemplate::new(401))
619            .expect(1)
620            .mount(&server)
621            .await;
622        wiremock::Mock::given(wiremock::matchers::method("GET"))
623            .and(wiremock::matchers::path("/test"))
624            .and(wiremock::matchers::header(
625                "Authorization",
626                "Bearer refreshed-token",
627            ))
628            .respond_with(wiremock::ResponseTemplate::new(200).set_body_bytes(b"ok".to_vec()))
629            .expect(1)
630            .mount(&server)
631            .await;
632
633        let resp = client
634            .get_bytes(&format!("{}/test", server.uri()))
635            .await
636            .unwrap();
637        assert_eq!(resp.status().as_u16(), 200);
638    }
639
640    #[tokio::test]
641    async fn get_bytes_propagates_network_errors() {
642        let client = DriveClient::new("http://127.0.0.1:1", &test_credentials()).unwrap();
643        let result = client.get_bytes("http://127.0.0.1:1/test").await;
644        assert!(result.is_err());
645    }
646
647    #[tokio::test]
648    async fn post_json_sends_body_and_bearer_auth() {
649        let server = wiremock::MockServer::start().await;
650        let client = client_with_bootstrapped_token(&server).await;
651        wiremock::Mock::given(wiremock::matchers::method("POST"))
652            .and(wiremock::matchers::path("/test"))
653            .and(wiremock::matchers::header(
654                "Authorization",
655                "Bearer bootstrap-token",
656            ))
657            .and(wiremock::matchers::body_json(serde_json::json!({"k": "v"})))
658            .respond_with(wiremock::ResponseTemplate::new(200))
659            .expect(1)
660            .mount(&server)
661            .await;
662
663        let resp = client
664            .post_json(
665                &format!("{}/test", server.uri()),
666                &serde_json::json!({"k": "v"}),
667            )
668            .await
669            .unwrap();
670        assert!(resp.status().is_success());
671    }
672
673    #[tokio::test]
674    async fn patch_json_sends_body_and_bearer_auth() {
675        let server = wiremock::MockServer::start().await;
676        let client = client_with_bootstrapped_token(&server).await;
677        wiremock::Mock::given(wiremock::matchers::method("PATCH"))
678            .and(wiremock::matchers::path("/test"))
679            .and(wiremock::matchers::header(
680                "Authorization",
681                "Bearer bootstrap-token",
682            ))
683            .and(wiremock::matchers::body_json(
684                serde_json::json!({"name": "new-name"}),
685            ))
686            .respond_with(wiremock::ResponseTemplate::new(200))
687            .expect(1)
688            .mount(&server)
689            .await;
690
691        let resp = client
692            .patch_json(
693                &format!("{}/test", server.uri()),
694                &serde_json::json!({"name": "new-name"}),
695            )
696            .await
697            .unwrap();
698        assert!(resp.status().is_success());
699    }
700
701    #[tokio::test]
702    async fn get_json_retries_on_429() {
703        let server = wiremock::MockServer::start().await;
704        let client = client_with_bootstrapped_token(&server).await;
705        wiremock::Mock::given(wiremock::matchers::method("GET"))
706            .and(wiremock::matchers::path("/test"))
707            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
708            .up_to_n_times(1)
709            .with_priority(1)
710            .mount(&server)
711            .await;
712        wiremock::Mock::given(wiremock::matchers::method("GET"))
713            .and(wiremock::matchers::path("/test"))
714            .respond_with(wiremock::ResponseTemplate::new(200))
715            .with_priority(2)
716            .mount(&server)
717            .await;
718
719        let resp = client
720            .get_json(&format!("{}/test", server.uri()))
721            .await
722            .unwrap();
723        assert_eq!(resp.status().as_u16(), 200);
724    }
725
726    #[tokio::test]
727    async fn get_json_retries_403_user_rate_limit_exceeded_then_succeeds() {
728        let server = wiremock::MockServer::start().await;
729        let client = client_with_bootstrapped_token(&server).await;
730        wiremock::Mock::given(wiremock::matchers::method("GET"))
731            .and(wiremock::matchers::path("/test"))
732            .respond_with(
733                wiremock::ResponseTemplate::new(403)
734                    .append_header("Retry-After", "0")
735                    .set_body_json(serde_json::json!({
736                        "error": {"message": "User Rate Limit Exceeded", "errors": [{"reason": "userRateLimitExceeded"}]}
737                    })),
738            )
739            .up_to_n_times(1)
740            .with_priority(1)
741            .mount(&server)
742            .await;
743        wiremock::Mock::given(wiremock::matchers::method("GET"))
744            .and(wiremock::matchers::path("/test"))
745            .respond_with(wiremock::ResponseTemplate::new(200))
746            .with_priority(2)
747            .mount(&server)
748            .await;
749
750        let resp = client
751            .get_json(&format!("{}/test", server.uri()))
752            .await
753            .unwrap();
754        assert_eq!(resp.status().as_u16(), 200);
755    }
756
757    #[tokio::test]
758    async fn get_json_does_not_retry_insufficient_permissions_403() {
759        let server = wiremock::MockServer::start().await;
760        let client = client_with_bootstrapped_token(&server).await;
761        wiremock::Mock::given(wiremock::matchers::method("GET"))
762            .and(wiremock::matchers::path("/test"))
763            .respond_with(
764                wiremock::ResponseTemplate::new(403).set_body_json(serde_json::json!({
765                    "error": {"message": "Insufficient Permission", "errors": [{"reason": "insufficientPermissions"}]}
766                })),
767            )
768            .expect(1)
769            .mount(&server)
770            .await;
771
772        let resp = client
773            .get_json(&format!("{}/test", server.uri()))
774            .await
775            .unwrap();
776        assert_eq!(resp.status().as_u16(), 403);
777    }
778
779    #[tokio::test]
780    async fn get_json_refreshes_and_retries_once_on_401() {
781        let server = wiremock::MockServer::start().await;
782        let client = client_with_bootstrapped_token(&server).await;
783        // The refresh endpoint issues a second, distinct token.
784        wiremock::Mock::given(wiremock::matchers::method("POST"))
785            .and(wiremock::matchers::path("/token"))
786            .respond_with(
787                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
788                    "access_token": "refreshed-token",
789                    "expires_in": 3600,
790                })),
791            )
792            .up_to_n_times(1)
793            .with_priority(2)
794            .mount(&server)
795            .await;
796        wiremock::Mock::given(wiremock::matchers::method("GET"))
797            .and(wiremock::matchers::path("/test"))
798            .and(wiremock::matchers::header(
799                "Authorization",
800                "Bearer bootstrap-token",
801            ))
802            .respond_with(wiremock::ResponseTemplate::new(401))
803            .expect(1)
804            .mount(&server)
805            .await;
806        wiremock::Mock::given(wiremock::matchers::method("GET"))
807            .and(wiremock::matchers::path("/test"))
808            .and(wiremock::matchers::header(
809                "Authorization",
810                "Bearer refreshed-token",
811            ))
812            .respond_with(wiremock::ResponseTemplate::new(200))
813            .expect(1)
814            .mount(&server)
815            .await;
816
817        let resp = client
818            .get_json(&format!("{}/test", server.uri()))
819            .await
820            .unwrap();
821        assert_eq!(resp.status().as_u16(), 200);
822    }
823
824    #[tokio::test]
825    async fn get_json_does_not_retry_a_second_time_on_persistent_401() {
826        let server = wiremock::MockServer::start().await;
827        let client = client_with_bootstrapped_token(&server).await;
828        wiremock::Mock::given(wiremock::matchers::method("POST"))
829            .and(wiremock::matchers::path("/token"))
830            .respond_with(
831                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
832                    "access_token": "still-rejected-token",
833                    "expires_in": 3600,
834                })),
835            )
836            .up_to_n_times(1)
837            .with_priority(2)
838            .mount(&server)
839            .await;
840        wiremock::Mock::given(wiremock::matchers::method("GET"))
841            .and(wiremock::matchers::path("/test"))
842            .respond_with(
843                wiremock::ResponseTemplate::new(401).set_body_string("still unauthorized"),
844            )
845            .expect(2)
846            .mount(&server)
847            .await;
848
849        let resp = client
850            .get_json(&format!("{}/test", server.uri()))
851            .await
852            .unwrap();
853        assert_eq!(resp.status().as_u16(), 401);
854    }
855
856    #[tokio::test]
857    async fn response_to_error_extracts_drive_message_and_reason() {
858        let server = wiremock::MockServer::start().await;
859        let client = client_with_bootstrapped_token(&server).await;
860        // `userRateLimitExceeded` is now retryable (`is_drive_quota_exceeded`),
861        // so without a zero-delay `Retry-After` this test would wait through
862        // the real exponential backoff before giving up.
863        wiremock::Mock::given(wiremock::matchers::method("GET"))
864            .and(wiremock::matchers::path("/test"))
865            .respond_with(
866                wiremock::ResponseTemplate::new(403)
867                    .append_header("Retry-After", "0")
868                    .set_body_json(serde_json::json!({
869                        "error": {
870                            "message": "User Rate Limit Exceeded",
871                            "errors": [{"reason": "userRateLimitExceeded"}],
872                        }
873                    })),
874            )
875            .mount(&server)
876            .await;
877
878        let resp = client
879            .get_json(&format!("{}/test", server.uri()))
880            .await
881            .unwrap();
882        let err = DriveClient::response_to_error(resp).await;
883        let msg = err.to_string();
884        assert!(msg.contains("User Rate Limit Exceeded"));
885        assert!(msg.contains("userRateLimitExceeded"));
886        assert_eq!(err.reason(), Some("userRateLimitExceeded"));
887    }
888
889    #[tokio::test]
890    async fn response_to_error_omits_reason_suffix_when_absent() {
891        let server = wiremock::MockServer::start().await;
892        let client = client_with_bootstrapped_token(&server).await;
893        wiremock::Mock::given(wiremock::matchers::method("GET"))
894            .and(wiremock::matchers::path("/test"))
895            .respond_with(
896                wiremock::ResponseTemplate::new(400).set_body_json(serde_json::json!({
897                    "error": {
898                        "message": "Invalid request",
899                    }
900                })),
901            )
902            .mount(&server)
903            .await;
904
905        let resp = client
906            .get_json(&format!("{}/test", server.uri()))
907            .await
908            .unwrap();
909        let err = DriveClient::response_to_error(resp).await;
910        let msg = err.to_string();
911        assert!(msg.contains("Invalid request"));
912        assert!(!msg.contains("reason:"));
913        assert_eq!(err.reason(), None);
914    }
915
916    #[tokio::test]
917    async fn response_to_error_falls_back_to_raw_body_when_not_drive_shaped() {
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(wiremock::ResponseTemplate::new(500).set_body_string("internal error"))
923            .mount(&server)
924            .await;
925
926        let resp = client
927            .get_json(&format!("{}/test", server.uri()))
928            .await
929            .unwrap();
930        let err = DriveClient::response_to_error(resp).await;
931        assert!(err.to_string().contains("internal error"));
932    }
933
934    #[tokio::test]
935    async fn get_json_propagates_network_errors() {
936        let client = DriveClient::new("http://127.0.0.1:1", &test_credentials()).unwrap();
937        let result = client.get_json("http://127.0.0.1:1/test").await;
938        assert!(result.is_err());
939    }
940
941    #[tokio::test]
942    async fn get_parsed_errors_on_malformed_json_response() {
943        let server = wiremock::MockServer::start().await;
944        let client = client_with_bootstrapped_token(&server).await;
945        wiremock::Mock::given(wiremock::matchers::method("GET"))
946            .and(wiremock::matchers::path("/test"))
947            .respond_with(wiremock::ResponseTemplate::new(200).set_body_string("not json"))
948            .mount(&server)
949            .await;
950
951        let result: Result<serde_json::Value> = client
952            .get_parsed(&format!("{}/test", server.uri()), "test context")
953            .await;
954        assert!(result.is_err());
955    }
956
957    #[tokio::test]
958    async fn get_parsed_errors_on_non_success_status_without_parsing_the_body() {
959        let server = wiremock::MockServer::start().await;
960        let client = client_with_bootstrapped_token(&server).await;
961        wiremock::Mock::given(wiremock::matchers::method("GET"))
962            .and(wiremock::matchers::path("/test"))
963            .respond_with(
964                wiremock::ResponseTemplate::new(404).set_body_json(serde_json::json!({
965                    "error": {"message": "File not found"}
966                })),
967            )
968            .mount(&server)
969            .await;
970
971        let result: Result<serde_json::Value> = client
972            .get_parsed(&format!("{}/test", server.uri()), "test context")
973            .await;
974        let err = result.unwrap_err();
975        assert!(err.to_string().contains("File not found"));
976    }
977}
omni_dev/drive/client.rs

omni_dev/drive/
client.rs
