Use AsyncMock and fixtures in co2signal tests (#104041)

homeassistant/components/co2signal/config_flow.py

@@ -120,18 +120,18 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         errors: dict[str, str] = {}
 
         session = async_get_clientsession(self.hass)
-        async with ElectricityMaps(token=data[CONF_API_KEY], session=session) as em:
-            try:
-                await fetch_latest_carbon_intensity(self.hass, em, data)
-            except InvalidToken:
-                errors["base"] = "invalid_auth"
-            except ElectricityMapsError:
-                errors["base"] = "unknown"
-            else:
-                return self.async_create_entry(
-                    title=get_extra_name(data) or "CO2 Signal",
-                    data=data,
-                )
+        em = ElectricityMaps(token=data[CONF_API_KEY], session=session)
+        try:
+            await fetch_latest_carbon_intensity(self.hass, em, data)
+        except InvalidToken:
+            errors["base"] = "invalid_auth"
+        except ElectricityMapsError:
+            errors["base"] = "unknown"
+        else:
+            return self.async_create_entry(
+                title=get_extra_name(data) or "CO2 Signal",
+                data=data,
+            )
 
         return self.async_show_form(
             step_id=step_id,

homeassistant/components/co2signal/coordinator.py

@@ -39,12 +39,11 @@ class CO2SignalCoordinator(DataUpdateCoordinator[CarbonIntensityResponse]):
     async def _async_update_data(self) -> CarbonIntensityResponse:
         """Fetch the latest data from the source."""
 
-        async with self.client as em:
-            try:
-                return await fetch_latest_carbon_intensity(
-                    self.hass, em, self.config_entry.data
-                )
-            except InvalidToken as err:
-                raise ConfigEntryError from err
-            except ElectricityMapsError as err:
-                raise UpdateFailed(str(err)) from err
+        try:
+            return await fetch_latest_carbon_intensity(
+                self.hass, self.client, self.config_entry.data
+            )
+        except InvalidToken as err:
+            raise ConfigEntryError from err
+        except ElectricityMapsError as err:
+            raise UpdateFailed(str(err)) from err

tests/components/co2signal/__init__.py

@@ -1,11 +1,18 @@
 """Tests for the CO2 Signal integration."""
+from aioelectricitymaps.models import (
+    CarbonIntensityData,
+    CarbonIntensityResponse,
+    CarbonIntensityUnit,
+)
 
-VALID_PAYLOAD = {
-    "status": "ok",
-    "countryCode": "FR",
-    "data": {
-        "carbonIntensity": 45.98623190095805,
-        "fossilFuelPercentage": 5.461182741937103,
-    },
-    "units": {"carbonIntensity": "gCO2eq/kWh"},
-}
+VALID_RESPONSE = CarbonIntensityResponse(
+    status="ok",
+    country_code="FR",
+    data=CarbonIntensityData(
+        carbon_intensity=45.98623190095805,
+        fossil_fuel_percentage=5.461182741937103,
+    ),
+    units=CarbonIntensityUnit(
+        carbon_intensity="gCO2eq/kWh",
+    ),
+)

tests/components/co2signal/conftest.py

@@ -0,0 +1,52 @@
+"""Fixtures for Electricity maps integration tests."""
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from homeassistant.components.co2signal import DOMAIN
+from homeassistant.const import CONF_API_KEY
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
+from tests.components.co2signal import VALID_RESPONSE
+
+
+@pytest.fixture(name="electricity_maps")
+def mock_electricity_maps() -> Generator[None, MagicMock, None]:
+    """Mock the ElectricityMaps client."""
+
+    with patch(
+        "homeassistant.components.co2signal.ElectricityMaps",
+        autospec=True,
+    ) as electricity_maps, patch(
+        "homeassistant.components.co2signal.config_flow.ElectricityMaps",
+        new=electricity_maps,
+    ):
+        client = electricity_maps.return_value
+        client.latest_carbon_intensity_by_coordinates.return_value = VALID_RESPONSE
+        client.latest_carbon_intensity_by_country_code.return_value = VALID_RESPONSE
+
+        yield client
+
+
+@pytest.fixture(name="config_entry")
+async def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry:
+    """Return a MockConfigEntry for testing."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_API_KEY: "api_key", "location": ""},
+        entry_id="904a74160aa6f335526706bee85dfb83",
+    )
+
+
+@pytest.fixture(name="setup_integration")
+async def mock_setup_integration(
+    hass: HomeAssistant, config_entry: MockConfigEntry, electricity_maps: AsyncMock
+) -> None:
+    """Fixture for setting up the component."""
+    config_entry.add_to_hass(hass)
+
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()

tests/components/co2signal/test_config_flow.py

@@ -1,5 +1,5 @@
 """Test the CO2 Signal config flow."""
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 from aioelectricitymaps.exceptions import (
     ElectricityMapsDecodeError,
@@ -13,9 +13,8 @@ from homeassistant.components.co2signal import DOMAIN, config_flow
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
-from . import VALID_PAYLOAD
-
 
+@pytest.mark.usefixtures("electricity_maps")
 async def test_form_home(hass: HomeAssistant) -> None:
     """Test we get the form."""
 
@@ -26,9 +25,6 @@ async def test_form_home(hass: HomeAssistant) -> None:
     assert result["errors"] is None
 
     with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ), patch(
         "homeassistant.components.co2signal.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
@@ -49,6 +45,7 @@ async def test_form_home(hass: HomeAssistant) -> None:
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+@pytest.mark.usefixtures("electricity_maps")
 async def test_form_coordinates(hass: HomeAssistant) -> None:
     """Test we get the form."""
 
@@ -68,9 +65,6 @@ async def test_form_coordinates(hass: HomeAssistant) -> None:
     assert result2["type"] == FlowResultType.FORM
 
     with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ), patch(
         "homeassistant.components.co2signal.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
@@ -93,6 +87,7 @@ async def test_form_coordinates(hass: HomeAssistant) -> None:
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+@pytest.mark.usefixtures("electricity_maps")
 async def test_form_country(hass: HomeAssistant) -> None:
     """Test we get the form."""
 
@@ -112,9 +107,6 @@ async def test_form_country(hass: HomeAssistant) -> None:
     assert result2["type"] == FlowResultType.FORM
 
     with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ), patch(
         "homeassistant.components.co2signal.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
@@ -151,39 +143,43 @@ async def test_form_country(hass: HomeAssistant) -> None:
         "json decode error",
     ],
 )
-async def test_form_error_handling(hass: HomeAssistant, side_effect, err_code) -> None:
+async def test_form_error_handling(
+    hass: HomeAssistant,
+    electricity_maps: AsyncMock,
+    side_effect: Exception,
+    err_code: str,
+) -> None:
     """Test we handle expected errors."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        side_effect=side_effect,
-    ):
-        result = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
-            {
-                "location": config_flow.TYPE_USE_HOME,
-                "api_key": "api_key",
-            },
-        )
+    electricity_maps.latest_carbon_intensity_by_coordinates.side_effect = side_effect
+    electricity_maps.latest_carbon_intensity_by_country_code.side_effect = side_effect
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "location": config_flow.TYPE_USE_HOME,
+            "api_key": "api_key",
+        },
+    )
 
     assert result["type"] == FlowResultType.FORM
     assert result["errors"] == {"base": err_code}
 
-    with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ):
-        result = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
-            {
-                "location": config_flow.TYPE_USE_HOME,
-                "api_key": "api_key",
-            },
-        )
-        await hass.async_block_till_done()
+    # reset mock and test if now succeeds
+    electricity_maps.latest_carbon_intensity_by_coordinates.side_effect = None
+    electricity_maps.latest_carbon_intensity_by_country_code.side_effect = None
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "location": config_flow.TYPE_USE_HOME,
+            "api_key": "api_key",
+        },
+    )
+    await hass.async_block_till_done()
 
     assert result["type"] == FlowResultType.CREATE_ENTRY
     assert result["title"] == "CO2 Signal"

tests/components/co2signal/test_diagnostics.py

@@ -1,38 +1,23 @@
 """Test the CO2Signal diagnostics."""
-from unittest.mock import patch
 
+import pytest
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.co2signal import DOMAIN
-from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
-from homeassistant.setup import async_setup_component
-
-from . import VALID_PAYLOAD
 
 from tests.common import MockConfigEntry
 from tests.components.diagnostics import get_diagnostics_for_config_entry
 from tests.typing import ClientSessionGenerator
 
 
+@pytest.mark.usefixtures("setup_integration")
 async def test_entry_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
+    config_entry: MockConfigEntry,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Test config entry diagnostics."""
-    config_entry = MockConfigEntry(
-        domain=DOMAIN,
-        data={CONF_API_KEY: "api_key", "location": ""},
-        entry_id="904a74160aa6f335526706bee85dfb83",
-    )
-    config_entry.add_to_hass(hass)
-    with patch(
-        "homeassistant.components.co2signal.coordinator.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ):
-        assert await async_setup_component(hass, DOMAIN, {})
-
     result = await get_diagnostics_for_config_entry(hass, hass_client, config_entry)
 
     assert result == snapshot