From 535e7dc70ba4ddbc68aa2aa4523a9fef63a9bdb7 Mon Sep 17 00:00:00 2001
From: Julio Merino <julio@meroh.net>
Date: Sun, 3 Dec 2023 08:17:07 -0800
Subject: [PATCH] Workaround Azure Maps 401 responses

Sometimes, Azure geolocation tests in pull requests fail with a 401,
saying that the Maps key is wrong.  However, this only happens to a
test within the same PR run and only randomly, so the configuration
is correct.

Try to workaround this with retries, which are also a good idea to
have because we are making the tests depend on an external service.
---
 geo/src/azure.rs | 45 ++++++++++++++++++++++++++++++---------------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/geo/src/azure.rs b/geo/src/azure.rs
index 1bb1b9e..cf67253 100644
--- a/geo/src/azure.rs
+++ b/geo/src/azure.rs
@@ -192,6 +192,7 @@ impl GeoLocator for AzureGeoLocator {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::time::Duration;
 
     #[test]
     pub fn test_azuregeolocatoroptions_from_env_all_present() {
@@ -223,26 +224,40 @@ mod tests {
         AzureGeoLocator::new(AzureGeoLocatorOptions::from_env("AZURE_MAPS").unwrap())
     }
 
+    /// Performs an Azure geolocation query with retries.
+    ///
+    /// I've observed that, sometimes, the queries randomly return a 401 Unauthorized error, but we
+    /// are passing the key correctly.  Retrying might help under this condition and others.
+    async fn geolocate(
+        geolocator: &AzureGeoLocator,
+        ip: &str,
+    ) -> GeoResult<Option<CountryIsoCode>> {
+        let ip = ip.parse().unwrap();
+
+        let mut retries = 5;
+        loop {
+            match geolocator.locate(&ip).await {
+                Ok(result) => return Ok(result),
+                Err(e) => {
+                    if retries == 0 {
+                        return Err(e);
+                    }
+                    retries -= 1;
+                    tokio::time::sleep(Duration::from_secs(1)).await;
+                }
+            }
+        }
+    }
+
     #[tokio::test]
     #[ignore = "Requires environment configuration and is expensive"]
     async fn test_ok() {
         let geolocator = setup();
-        assert_eq!(
-            "ES",
-            geolocator.locate(&"212.170.36.79".parse().unwrap()).await.unwrap().unwrap().as_str()
-        );
-        assert_eq!(
-            "IE",
-            geolocator.locate(&"185.2.66.42".parse().unwrap()).await.unwrap().unwrap().as_str()
-        );
+        assert_eq!("ES", geolocate(&geolocator, "212.170.36.79").await.unwrap().unwrap().as_str());
+        assert_eq!("IE", geolocate(&geolocator, "185.2.66.42").await.unwrap().unwrap().as_str());
         assert_eq!(
             "US",
-            geolocator
-                .locate(&"2001:4898:80e8:3c::".parse().unwrap())
-                .await
-                .unwrap()
-                .unwrap()
-                .as_str()
+            geolocate(&geolocator, "2001:4898:80e8:3c::").await.unwrap().unwrap().as_str()
         );
     }
 
@@ -250,6 +265,6 @@ mod tests {
     #[ignore = "Requires environment configuration and is expensive"]
     async fn test_missing() {
         let geolocator = setup();
-        assert_eq!(None, geolocator.locate(&"198.18.0.1".parse().unwrap()).await.unwrap());
+        assert_eq!(None, geolocate(&geolocator, "198.18.0.1").await.unwrap());
     }
 }