Use AsyncMock and fixtures in co2signal tests (#104041)

homeassistant/components/co2signal/config_flow.py

@@ -120,7 +120,7 @@ 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:
+        em = ElectricityMaps(token=data[CONF_API_KEY], session=session)
         try:
             await fetch_latest_carbon_intensity(self.hass, em, data)
         except InvalidToken:

homeassistant/components/co2signal/coordinator.py

@@ -39,10 +39,9 @@ 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
+                self.hass, self.client, self.config_entry.data
             )
         except InvalidToken as err:
             raise ConfigEntryError 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,16 +143,20 @@ 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,
-    ):
+    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"],
         {
@@ -172,10 +168,10 @@ async def test_form_error_handling(hass: HomeAssistant, side_effect, err_code) -
     assert result["type"] == FlowResultType.FORM
     assert result["errors"] == {"base": err_code}
 
-    with patch(
-        "homeassistant.components.co2signal.config_flow.ElectricityMaps._get",
-        return_value=VALID_PAYLOAD,
-    ):
+    # 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"],
         {

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