Fix unit testing: Pass auth-token to restful API

to avoid rate limit Signed-off-by: Jiahao XU <[email protected]>
cargo-bins · Jun 12, 2024 · 1f363fe · 1f363fe
1 parent 54149a0
commit 1f363fe
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 15 deletions.
diff --git a/crates/binstalk-git-repo-api/src/gh_api_client.rs b/crates/binstalk-git-repo-api/src/gh_api_client.rs
@@ -131,6 +131,8 @@ struct Inner {
 
     auth_token: Option<CompactString>,
     is_auth_token_valid: AtomicBool,
+
+    only_use_restful_api: AtomicBool,
 }
 
 /// Github API client for querying whether a release artifact exitsts.
@@ -147,9 +149,16 @@ impl GhApiClient {
 
             auth_token,
             is_auth_token_valid: AtomicBool::new(true),
+
+            only_use_restful_api: AtomicBool::new(false),
         }))
     }
 
+    /// If you don't want to use GitHub GraphQL API for whatever reason, call this.
+    pub fn set_only_use_restful_api(&self) {
+        self.0.only_use_restful_api.store(true, Relaxed);
+    }
+
     pub fn remote_client(&self) -> &remote::Client {
         &self.0.client
     }
@@ -193,22 +202,24 @@ impl GhApiClient {
     ) -> Result<U, GhApiError>
     where
         GraphQLFn: Fn(&remote::Client, &T, &str) -> GraphQLFut,
-        RestfulFn: Fn(&remote::Client, &T) -> RestfulFut,
+        RestfulFn: Fn(&remote::Client, &T, Option<&str>) -> RestfulFut,
         GraphQLFut: Future<Output = Result<U, GhApiError>> + Send + Sync + 'static,
         RestfulFut: Future<Output = Result<U, GhApiError>> + Send + Sync + 'static,
     {
         self.check_retry_after()?;
 
-        if let Some(auth_token) = self.get_auth_token() {
-            match graphql_func(&self.0.client, data, auth_token).await {
-                Err(GhApiError::Unauthorized) => {
-                    self.0.is_auth_token_valid.store(false, Relaxed);
+        if !self.0.only_use_restful_api.load(Relaxed) {
+            if let Some(auth_token) = self.get_auth_token() {
+                match graphql_func(&self.0.client, data, auth_token).await {
+                    Err(GhApiError::Unauthorized) => {
+                        self.0.is_auth_token_valid.store(false, Relaxed);
+                    }
+                    res => return res.map_err(|err| err.context("GraphQL API")),
                 }
-                res => return res.map_err(|err| err.context("GraphQL API")),
             }
         }
 
-        restful_func(&self.0.client, data)
+        restful_func(&self.0.client, data, self.get_auth_token())
             .await
             .map_err(|err| err.context("Restful API"))
     }
@@ -369,6 +380,7 @@ mod test {
             tag: CompactString::new_inline("cargo-audit/v0.17.6"),
         };
 
+        #[allow(unused)]
         pub(super) const ARTIFACTS: &[&str] = &[
             "cargo-audit-aarch64-unknown-linux-gnu-v0.17.6.tgz",
             "cargo-audit-armv7-unknown-linux-gnueabihf-v0.17.6.tgz",
@@ -512,10 +524,17 @@ mod test {
     fn create_client() -> Vec<GhApiClient> {
         let client = create_remote_client();
 
-        let mut gh_clients = vec![GhApiClient::new(client.clone(), None)];
+        let auth_token = env::var("CI_UNIT_TEST_GITHUB_TOKEN")
+            .ok()
+            .map(CompactString::from);
+
+        let gh_client = GhApiClient::new(client.clone(), auth_token.clone());
+        gh_client.set_only_use_restful_api();
+
+        let mut gh_clients = vec![gh_client];
 
-        if let Ok(token) = env::var("CI_UNIT_TEST_GITHUB_TOKEN") {
-            gh_clients.push(GhApiClient::new(client, Some(token.into())));
+        if auth_token.is_some() {
+            gh_clients.push(GhApiClient::new(client, auth_token));
         }
 
         gh_clients

diff --git a/crates/binstalk-git-repo-api/src/gh_api_client/common.rs b/crates/binstalk-git-repo-api/src/gh_api_client/common.rs
@@ -38,6 +38,7 @@ fn get_api_endpoint() -> &'static Url {
 pub(super) fn issue_restful_api<T>(
     client: &remote::Client,
     path: &[&str],
+    auth_token: Option<&str>,
 ) -> impl Future<Output = Result<T, GhApiError>> + Send + Sync + 'static
 where
     T: DeserializeOwned,
@@ -50,11 +51,16 @@ where
 
     debug!("Getting restful API: {url}");
 
-    let future = client
+    let mut request_builder = client
         .get(url)
         .header("Accept", "application/vnd.github+json")
-        .header("X-GitHub-Api-Version", "2022-11-28")
-        .send(false);
+        .header("X-GitHub-Api-Version", "2022-11-28");
+
+    if let Some(auth_token) = auth_token {
+        request_builder = request_builder.bearer_auth(&auth_token);
+    }
+
+    let future = request_builder.send(false);
 
     async move {
         let response = check_http_status_and_header(future.await?)?;

diff --git a/crates/binstalk-git-repo-api/src/gh_api_client/release_artifacts.rs b/crates/binstalk-git-repo-api/src/gh_api_client/release_artifacts.rs
@@ -73,8 +73,13 @@ pub(super) fn fetch_release_artifacts_restful_api(
         repo: GhRepo { owner, repo },
         tag,
     }: &GhRelease,
+    auth_token: Option<&str>,
 ) -> impl Future<Output = Result<Artifacts, GhApiError>> + Send + Sync + 'static {
-    issue_restful_api(client, &["repos", owner, repo, "releases", "tags", tag])
+    issue_restful_api(
+        client,
+        &["repos", owner, repo, "releases", "tags", tag],
+        auth_token,
+    )
 }
 
 #[derive(Debug, Deserialize)]

diff --git a/crates/binstalk-git-repo-api/src/gh_api_client/repo_info.rs b/crates/binstalk-git-repo-api/src/gh_api_client/repo_info.rs
@@ -44,8 +44,9 @@ impl RepoInfo {
 pub(super) fn fetch_repo_info_restful_api(
     client: &remote::Client,
     GhRepo { owner, repo }: &GhRepo,
+    auth_token: Option<&str>,
 ) -> impl Future<Output = Result<Option<RepoInfo>, GhApiError>> + Send + Sync + 'static {
-    issue_restful_api(client, &["repos", owner, repo])
+    issue_restful_api(client, &["repos", owner, repo], auth_token)
 }
 
 #[derive(Debug, Deserialize)]