Merge branch 'write-relay-selection-tests-for-relay-ip-overrides-des-579'

1 files changed, 12 insertions, 3 deletions

diff --git a/test/test-rpc/src/net.rs b/test/test-rpc/src/net.rs
index 77aa5c938a..124b4f9ad3 100644
--- a/test/test-rpc/src/net.rs
+++ b/test/test-rpc/src/net.rs
@@ -2,7 +2,7 @@ use futures::channel::oneshot;
 use hyper::{Client, Uri};
 use once_cell::sync::Lazy;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
-use std::net::SocketAddr;
+use std::{net::SocketAddr, time::Duration};
 use tokio_rustls::rustls::ClientConfig;
 
 use crate::{AmIMullvad, Error};
@@ -72,10 +72,10 @@ impl Drop for SockHandle {
     }
 }
 
-pub async fn geoip_lookup(mullvad_host: String) -> Result<AmIMullvad, Error> {
+pub async fn geoip_lookup(mullvad_host: String, timeout: Duration) -> Result<AmIMullvad, Error> {
     let uri = Uri::try_from(format!("https://ipv4.am.i.{mullvad_host}/json"))
         .map_err(|_| Error::InvalidUrl)?;
-    http_get(uri).await
+    http_get_with_timeout(uri, timeout).await
 }
 
 pub async fn http_get<T: DeserializeOwned>(url: Uri) -> Result<T, Error> {
@@ -106,6 +106,15 @@ pub async fn http_get<T: DeserializeOwned>(url: Uri) -> Result<T, Error> {
     })
 }
 
+pub async fn http_get_with_timeout<T: DeserializeOwned>(
+    url: Uri,
+    timeout: Duration,
+) -> Result<T, Error> {
+    tokio::time::timeout(timeout, http_get(url))
+        .await
+        .map_err(|_| Error::HttpRequest("Request timed out".into()))?
+}
+
 fn read_cert_store() -> tokio_rustls::rustls::RootCertStore {
     let mut cert_store = tokio_rustls::rustls::RootCertStore::empty();