Speedup binstalk-git-repo-api unit testing (#1731)

* Share `remote::Client` between tests when in cargo-test Signed-off-by: Jiahao XU <[email protected]> * Run rate-limited binstalk-git-repo-api test in serial Signed-off-by: Jiahao XU <[email protected]> * Update client-side rate limit to 1 request per 300ms Signed-off-by: Jiahao XU <[email protected]> * Reduce test cases in `rate_limited_test_has_release_artifact_and_download_artifacts` Signed-off-by: Jiahao XU <[email protected]> * Add `cargo-nextest` that I forgot Signed-off-by: Jiahao XU <[email protected]> * Fix unit testing: Pass auth-token to restful API to avoid rate limit Signed-off-by: Jiahao XU <[email protected]> --------- Signed-off-by: Jiahao XU <[email protected]>
cargo-bins · Jun 12, 2024 · 6622bf1 · 6622bf1
1 parent 4d7a91a
commit 6622bf1
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 35 deletions.
diff --git a/.config/nextest.toml b/.config/nextest.toml
@@ -0,0 +1,6 @@
+[test-groups]
+rate-limited = { max-threads = 1 }
+
+[[profile.default.overrides]]
+filter = 'test(rate_limited::)'
+test-group = 'rate-limited'
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/binstalk-git-repo-api/Cargo.toml b/crates/binstalk-git-repo-api/Cargo.toml
@@ -26,3 +26,4 @@ url = "2.3.1"
 [dev-dependencies]
 binstalk-downloader = { version = "0.11.1", path = "../binstalk-downloader" }
 tracing-subscriber = "0.3"
+once_cell = "1"
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"))
     }
@@ -313,6 +324,7 @@ impl GhApiClient {
 mod test {
     use super::*;
     use compact_str::{CompactString, ToCompactString};
+    use once_cell::sync::OnceCell;
     use std::{env, num::NonZeroU16, time::Duration};
     use tokio::time::sleep;
     use tracing::subscriber::set_global_default;
@@ -368,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",
@@ -490,32 +503,45 @@ mod test {
     }
 
     fn create_remote_client() -> remote::Client {
-        remote::Client::new(
-            concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION")),
-            None,
-            NonZeroU16::new(200).unwrap(),
-            1.try_into().unwrap(),
-            [],
-        )
-        .unwrap()
+        static CLIENT: OnceCell<remote::Client> = OnceCell::new();
+
+        CLIENT
+            .get_or_init(|| {
+                remote::Client::new(
+                    concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION")),
+                    None,
+                    NonZeroU16::new(300).unwrap(),
+                    1.try_into().unwrap(),
+                    [],
+                )
+                .unwrap()
+            })
+            .clone()
     }
 
     /// Mark this as an async fn so that you won't accidentally use it in
     /// sync context.
     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();
 
-        if let Ok(token) = env::var("CI_UNIT_TEST_GITHUB_TOKEN") {
-            gh_clients.push(GhApiClient::new(client, Some(token.into())));
+        let mut gh_clients = vec![gh_client];
+
+        if auth_token.is_some() {
+            gh_clients.push(GhApiClient::new(client, auth_token));
         }
 
         gh_clients
     }
 
     #[tokio::test]
-    async fn test_get_repo_info() {
+    async fn rate_limited_test_get_repo_info() {
         const PUBLIC_REPOS: [GhRepo; 1] = [GhRepo {
             owner: CompactString::new_inline("cargo-bins"),
             repo: CompactString::new_inline("cargo-binstall"),
@@ -575,17 +601,11 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_has_release_artifact_and_download_artifacts() {
-        const RELEASES: [(GhRelease, &[&str]); 2] = [
-            (
-                cargo_binstall_v0_20_1::RELEASE,
-                cargo_binstall_v0_20_1::ARTIFACTS,
-            ),
-            (
-                cargo_audit_v_0_17_6::RELEASE,
-                cargo_audit_v_0_17_6::ARTIFACTS,
-            ),
-        ];
+    async fn rate_limited_test_has_release_artifact_and_download_artifacts() {
+        const RELEASES: [(GhRelease, &[&str]); 1] = [(
+            cargo_binstall_v0_20_1::RELEASE,
+            cargo_binstall_v0_20_1::ARTIFACTS,
+        )];
         const NON_EXISTENT_RELEASES: [GhRelease; 1] = [GhRelease {
             repo: GhRepo {
                 owner: CompactString::new_inline("cargo-bins"),

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)]