Fix retries (#1573)

lychee-bin/tests/cli.rs

@@ -1880,4 +1880,28 @@ mod cli {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_retry() -> Result<()> {
+        let mock_server = wiremock::MockServer::start().await;
+
+        wiremock::Mock::given(wiremock::matchers::method("GET"))
+            .respond_with(ResponseTemplate::new(429))
+            .up_to_n_times(1)
+            .mount(&mock_server)
+            .await;
+
+        wiremock::Mock::given(wiremock::matchers::method("GET"))
+            .respond_with(ResponseTemplate::new(200))
+            .mount(&mock_server)
+            .await;
+
+        let mut cmd = main_command();
+        cmd.arg("-")
+            .write_stdin(mock_server.uri())
+            .assert()
+            .success();
+
+        Ok(())
+    }
 }

lychee-lib/src/retry.rs

@@ -16,11 +16,11 @@ pub(crate) trait RetryExt {
     fn should_retry(&self) -> bool;
 }
 
-impl RetryExt for reqwest::Response {
+impl RetryExt for reqwest::StatusCode {
     /// Try to map a `reqwest` response into `Retryable`.
     #[allow(clippy::if_same_then_else)]
     fn should_retry(&self) -> bool {
-        let status = self.status();
+        let status = *self;
         if status.is_server_error() {
             true
         } else if status.is_client_error()
@@ -71,6 +71,8 @@ impl RetryExt for reqwest::Error {
             } else {
                 false
             }
+        } else if let Some(status) = self.status() {
+            status.should_retry()
         } else {
             // We omit checking if error.is_status() since we check that already.
             // However, if Response::error_for_status is used the status will still
@@ -148,3 +150,18 @@ fn get_source_error_type<T: std::error::Error + 'static>(
     }
     None
 }
+
+#[cfg(test)]
+mod tests {
+    use http::StatusCode;
+
+    use super::RetryExt;
+
+    #[test]
+    fn test_should_retry() {
+        assert!(StatusCode::REQUEST_TIMEOUT.should_retry());
+        assert!(StatusCode::TOO_MANY_REQUESTS.should_retry());
+        assert!(!StatusCode::FORBIDDEN.should_retry());
+        assert!(StatusCode::INTERNAL_SERVER_ERROR.should_retry());
+    }
+}