Refactor Influx logic to reduce V1 vs V2 code paths (#37232)

* refactoring to share logic and sensor startup error test

* Added handling for V1 InfluxDBServerError to start-up and runtime and test for it

* Added InfluxDBServerError test to sensor setup tests

* Raising PlatformNotReady exception from sensor for setup failure

* Proper testing of PlatformNotReady error

tests/components/influxdb/test_sensor.py

@@ -1,8 +1,9 @@
 """The tests for the InfluxDB sensor."""
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Dict, List, Type
 
-from influxdb.exceptions import InfluxDBClientError
+from influxdb.exceptions import InfluxDBClientError, InfluxDBServerError
 from influxdb_client.rest import ApiException
 import pytest
 from voluptuous import Invalid
@@ -18,12 +19,15 @@ from homeassistant.components.influxdb.sensor import PLATFORM_SCHEMA
 import homeassistant.components.sensor as sensor
 from homeassistant.const import STATE_UNKNOWN
 from homeassistant.setup import async_setup_component
+from homeassistant.util import dt as dt_util
 
 from tests.async_mock import MagicMock, patch
+from tests.common import async_fire_time_changed
 
 INFLUXDB_PATH = "homeassistant.components.influxdb"
-INFLUXDB_CLIENT_PATH = f"{INFLUXDB_PATH}.sensor.InfluxDBClient"
+INFLUXDB_CLIENT_PATH = f"{INFLUXDB_PATH}.InfluxDBClient"
 INFLUXDB_SENSOR_PATH = f"{INFLUXDB_PATH}.sensor"
+PLATFORM_NOT_READY_BASE_WAIT_TIME = 30
 
 BASE_V1_CONFIG = {}
 BASE_V2_CONFIG = {
@@ -137,6 +141,8 @@ def _set_query_mock_v1(mock_influx_client, return_value=None, side_effect=None):
 
         query_api.side_effect = get_return_value
 
+    return query_api
+
 
 def _set_query_mock_v2(mock_influx_client, return_value=None, side_effect=None):
     """Set return value or side effect for the V2 client."""
@@ -149,6 +155,8 @@ def _set_query_mock_v2(mock_influx_client, return_value=None, side_effect=None):
 
         query_api.return_value = return_value
 
+    return query_api
+
 
 async def _setup(hass, config_ext, queries, expected_sensors):
     """Create client and test expected sensors."""
@@ -451,3 +459,79 @@ async def test_error_rendering_template(
     assert (
         len([record for record in caplog.records if record.levelname == "ERROR"]) == 1
     )
+
+
+@pytest.mark.parametrize(
+    "mock_client, config_ext, queries, set_query_mock, test_exception, make_resultset",
+    [
+        (
+            DEFAULT_API_VERSION,
+            BASE_V1_CONFIG,
+            BASE_V1_QUERY,
+            _set_query_mock_v1,
+            OSError("fail"),
+            _make_v1_resultset,
+        ),
+        (
+            DEFAULT_API_VERSION,
+            BASE_V1_CONFIG,
+            BASE_V1_QUERY,
+            _set_query_mock_v1,
+            InfluxDBClientError("fail"),
+            _make_v1_resultset,
+        ),
+        (
+            DEFAULT_API_VERSION,
+            BASE_V1_CONFIG,
+            BASE_V1_QUERY,
+            _set_query_mock_v1,
+            InfluxDBServerError("fail"),
+            _make_v1_resultset,
+        ),
+        (
+            API_VERSION_2,
+            BASE_V2_CONFIG,
+            BASE_V2_QUERY,
+            _set_query_mock_v2,
+            OSError("fail"),
+            _make_v2_resultset,
+        ),
+        (
+            API_VERSION_2,
+            BASE_V2_CONFIG,
+            BASE_V2_QUERY,
+            _set_query_mock_v2,
+            ApiException(),
+            _make_v2_resultset,
+        ),
+    ],
+    indirect=["mock_client"],
+)
+async def test_connection_error_at_startup(
+    hass,
+    caplog,
+    mock_client,
+    config_ext,
+    queries,
+    set_query_mock,
+    test_exception,
+    make_resultset,
+):
+    """Test behavior of sensor when influx returns error."""
+    query_api = set_query_mock(mock_client, side_effect=test_exception)
+    expected_sensor = "sensor.test"
+
+    # Test sensor is not setup first time due to connection error
+    await _setup(hass, config_ext, queries, [])
+    assert hass.states.get(expected_sensor) is None
+    assert (
+        len([record for record in caplog.records if record.levelname == "ERROR"]) == 1
+    )
+
+    # Stop throwing exception and advance time to test setup succeeds
+    query_api.reset_mock(side_effect=True)
+    set_query_mock(mock_client, return_value=make_resultset(42))
+    new_time = dt_util.utcnow() + timedelta(seconds=PLATFORM_NOT_READY_BASE_WAIT_TIME)
+    async_fire_time_changed(hass, new_time)
+    await hass.async_block_till_done()
+    assert hass.states.get(expected_sensor) is not None