Add test for midpoint

1 files changed, 58 insertions, 2 deletions

diff --git a/mullvad-types/src/location.rs b/mullvad-types/src/location.rs
index 9bf46cb586..2a98df4b8e 100644
--- a/mullvad-types/src/location.rs
+++ b/mullvad-types/src/location.rs
@@ -31,7 +31,7 @@ impl Location {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct Coordinates {
     pub latitude: f64,
     pub longitude: f64,
@@ -61,10 +61,16 @@ impl Coordinates {
     ///
     /// https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates
     pub fn midpoint(locations: &[Location]) -> Self {
+        Self::midpoint_inner(locations.iter().map(Coordinates::from))
+    }
+
+    fn midpoint_inner(locations: impl std::iter::Iterator<Item = Coordinates>) -> Self {
         let mut x = 0f64;
         let mut y = 0f64;
         let mut z = 0f64;
 
+        let mut count = 0;
+
         for location in locations {
             let cos_lat = location.latitude.to_radians().cos();
             let sin_lat = location.latitude.to_radians().sin();
@@ -73,8 +79,9 @@ impl Coordinates {
             x += cos_lat * cos_lon;
             y += cos_lat * sin_lon;
             z += sin_lat;
+            count += 1;
         }
-        let inv_total_weight = 1f64 / (locations.len() as f64);
+        let inv_total_weight = 1f64 / (count as f64);
         x *= inv_total_weight;
         y *= inv_total_weight;
         z *= inv_total_weight;
@@ -169,6 +176,16 @@ impl From<AmIMullvad> for GeoIpLocation {
 
 #[cfg(test)]
 mod tests {
+    use super::Coordinates;
+
+    impl Coordinates {
+        fn equal(&self, other: Coordinates) -> bool {
+            const EPS: f64 = 0.1;
+            (self.latitude - other.latitude).abs() < EPS
+                && (self.longitude - other.longitude).abs() < EPS
+        }
+    }
+
     #[test]
     fn test_haversine_dist_deg() {
         use super::haversine_dist_deg;
@@ -188,4 +205,43 @@ mod tests {
             111.22634257109495
         );
     }
+
+    #[test]
+    fn test_midpoint() {
+        assert!(Coordinates::midpoint_inner(
+            [
+                Coordinates {
+                    latitude: 0.0,
+                    longitude: 90.0,
+                },
+                Coordinates {
+                    latitude: 90.0,
+                    longitude: 0.0,
+                },
+            ]
+            .into_iter()
+        )
+        .equal(Coordinates {
+            latitude: 45.0,
+            longitude: 90.0,
+        }));
+
+        assert!(Coordinates::midpoint_inner(
+            [
+                Coordinates {
+                    latitude: -20.0,
+                    longitude: 90.0,
+                },
+                Coordinates {
+                    latitude: -20.0,
+                    longitude: -90.0,
+                },
+            ]
+            .into_iter()
+        )
+        .equal(Coordinates {
+            latitude: -90.0,
+            longitude: 0.0,
+        }));
+    }
 }