Fix tests

Airthings · Jan 1, 2025 · 2247a44 · 2247a44
1 parent 0f46817
commit 2247a44
Showing 1 changed file with 16 additions and 11 deletions.
diff --git a/tests/test_device_type.py b/tests/test_device_type.py
@@ -9,12 +9,14 @@ def test_device_type():
     assert AirthingsDeviceType.WAVE_PLUS.product_name == "Wave Plus"
     assert AirthingsDeviceType.WAVE_RADON.product_name == "Wave Radon"
     assert AirthingsDeviceType.WAVE_GEN_1.product_name == "Wave Gen 1"
-    assert AirthingsDeviceType.WAVE_ENHANCE.product_name == "Wave Enhance"
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.product_name == "Wave Enhance"
+    assert AirthingsDeviceType.WAVE_ENHANCE_US.product_name == "Wave Enhance"
     assert AirthingsDeviceType("2900") == AirthingsDeviceType.WAVE_GEN_1
     assert AirthingsDeviceType("2920") == AirthingsDeviceType.WAVE_MINI
     assert AirthingsDeviceType("2930") == AirthingsDeviceType.WAVE_PLUS
     assert AirthingsDeviceType("2950") == AirthingsDeviceType.WAVE_RADON
-    assert AirthingsDeviceType("3210") == AirthingsDeviceType.WAVE_ENHANCE
+    assert AirthingsDeviceType("3210") == AirthingsDeviceType.WAVE_ENHANCE_EU
+    assert AirthingsDeviceType("3220") == AirthingsDeviceType.WAVE_ENHANCE_US
     with pytest.raises(ValueError):
         AirthingsDeviceType("1234")
 
@@ -23,7 +25,10 @@ def test_device_type():
     assert AirthingsDeviceType.from_raw_value("2930") == AirthingsDeviceType.WAVE_PLUS
     assert AirthingsDeviceType.from_raw_value("2950") == AirthingsDeviceType.WAVE_RADON
     assert (
-        AirthingsDeviceType.from_raw_value("3210") == AirthingsDeviceType.WAVE_ENHANCE
+        AirthingsDeviceType.from_raw_value("3210") == AirthingsDeviceType.WAVE_ENHANCE_EU
+    )
+    assert (
+        AirthingsDeviceType.from_raw_value("3220") == AirthingsDeviceType.WAVE_ENHANCE_US
     )
 
     unknown_device = AirthingsDeviceType.from_raw_value("1234")
@@ -76,11 +81,11 @@ def test_battery_calculation():
     assert AirthingsDeviceType.WAVE_GEN_1.battery_percentage(2.1) == 0
     assert AirthingsDeviceType.WAVE_GEN_1.battery_percentage(2.0) == 0
 
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(3.2) == 100
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(3.0) == 100
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.8) == 81
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.6) == 53
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.5) == 28
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.2) == 5
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.1) == 0
-    assert AirthingsDeviceType.WAVE_ENHANCE.battery_percentage(2.0) == 0
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(3.2) == 100
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(3.0) == 100
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.8) == 81
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.6) == 53
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.5) == 28
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.2) == 5
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.1) == 0
+    assert AirthingsDeviceType.WAVE_ENHANCE_EU.battery_percentage(2.0) == 0