Merge 2aacef5 into 5e6c40b

flexmeasures/data/tests/test_sensor_queries.py

@@ -1,25 +1,43 @@
+import pytest
+
 from flexmeasures.data.models.time_series import Sensor
 
 
-def test_closest_sensor(run_as_cli, add_nearby_weather_sensors):
+@pytest.mark.parametrize("n", [1, 3])
+def test_closest_sensor(run_as_cli, add_nearby_weather_sensors, n):
     """Check that the closest temperature sensor to our wind sensor returns
     the one that is on the same spot as the wind sensor itself.
     (That's where we set it up in our conftest.)
-    And check that the 2nd and 3rd closest are the farther temperature sensors we set up.
+    And check that the 2nd and 3rd closest are the farther temperature sensors we set up, in the case of 3 sensors.
     """
     wind_sensor = add_nearby_weather_sensors["wind"]
-    closest_sensors = Sensor.find_closest(
+    closest_sensor_or_sensors = Sensor.find_closest(
         generic_asset_type_name=wind_sensor.generic_asset.generic_asset_type.name,
-        n=3,
+        n=n,
         sensor_name="temperature",
         latitude=wind_sensor.generic_asset.latitude,
         longitude=wind_sensor.generic_asset.longitude,
     )
-    assert closest_sensors[0].location == wind_sensor.generic_asset.location
-    assert closest_sensors[1] == add_nearby_weather_sensors["farther_temperature"]
-    assert closest_sensors[2] == add_nearby_weather_sensors["even_farther_temperature"]
-    for sensor in closest_sensors:
+    if n == 1:
+        assert closest_sensor_or_sensors.location == wind_sensor.generic_asset.location
         assert (
-            sensor.generic_asset.generic_asset_type.name
+            closest_sensor_or_sensors.generic_asset.generic_asset_type.name
             == wind_sensor.generic_asset.generic_asset_type.name
         )
+    elif n == 3:
+        assert (
+            closest_sensor_or_sensors[0].location == wind_sensor.generic_asset.location
+        )
+        assert (
+            closest_sensor_or_sensors[1]
+            == add_nearby_weather_sensors["farther_temperature"]
+        )
+        assert (
+            closest_sensor_or_sensors[2]
+            == add_nearby_weather_sensors["even_farther_temperature"]
+        )
+        for sensor in closest_sensor_or_sensors:
+            assert (
+                sensor.generic_asset.generic_asset_type.name
+                == wind_sensor.generic_asset.generic_asset_type.name
+            )