Switch to fixtures for influxdb mocking in data api tests

metobsapi/tests/test_data_api.py

+import contextlib
 import json
+import random
+from datetime import datetime, timedelta
+from typing import Callable, ContextManager, Iterable, Iterator, TypeAlias
 from unittest import mock
 
 import pytest
+from influxdb.resultset import ResultSet
 
+QueryResult: TypeAlias = ResultSet | list[ResultSet]
 
-def fake_data(interval, symbols, num_vals, single_result=False):
-    import random
-    from datetime import datetime, timedelta
 
-    from influxdb.resultset import ResultSet
+@pytest.fixture
+def mock_influxdb_query() -> Callable[[QueryResult], ContextManager[mock.Mock]]:
+    @contextlib.contextmanager
+    def _mock_influxdb_query_with_fake_data(fake_result_data: QueryResult) -> Iterator[mock.Mock]:
+        with mock.patch("metobsapi.data_api.query") as query_func:
+            query_func.return_value = fake_result_data
+            yield query_func
+
+    return _mock_influxdb_query_with_fake_data
 
+
+@pytest.fixture
+def influxdb_air_temp_9_values(mock_influxdb_query) -> Iterable[None]:
+    fake_result = _fake_data("1m", {("aoss", "tower"): ["time", "air_temp"]}, 9)
+    with mock_influxdb_query(fake_result):
+        yield
+
+
+@pytest.fixture
+def influxdb_3_symbols_9_values(mock_influxdb_query) -> Iterable[None]:
+    fake_result = _fake_data("1m", {("aoss", "tower"): ["time", "air_temp", "rel_hum", "wind_speed"]}, 9)
+    with mock_influxdb_query(fake_result):
+        yield
+
+
+@pytest.fixture
+def influxdb_wind_fields_9_values(mock_influxdb_query) -> Iterable[None]:
+    fake_result = _fake_data(
+        "1m", {("aoss", "tower"): ["time", "wind_speed", "wind_direction", "wind_east", "wind_north"]}, 9
+    )
+    with mock_influxdb_query(fake_result):
+        yield
+
+
+def _fake_data(interval, symbols, num_vals, single_result=False):
     now = datetime(2017, 3, 5, 19, 0, 0)
     t_format = "%Y-%m-%dT%H:%M:%SZ"
     measurement_name = "metobs_" + interval
@@ -133,10 +169,7 @@ class TestDataAPI:
         assert res["code"] == 413
         assert res["status"] == "fail"
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_shorthand_one_symbol_json_row(self, query_func, client):
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp"]}, 9)
-        query_func.return_value = r
+    def test_shorthand_one_symbol_json_row(self, client, influxdb_air_temp_9_values):
         # row should be the default
         res = client.get("/api/data.json?site=aoss&inst=tower&symbols=air_temp&begin=-00:10:00")
         res = json.loads(res.data.decode())
@@ -147,10 +180,7 @@ class TestDataAPI:
         assert len(res["results"]["data"]) == 9
         assert len(res["results"]["data"][0]) == 1
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_shorthand_one_symbol_json_column(self, query_func, client):
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp"]}, 9)
-        query_func.return_value = r
+    def test_shorthand_one_symbol_json_column(self, client, influxdb_air_temp_9_values):
         res = client.get("/api/data.json?site=aoss&inst=tower&symbols=air_temp&begin=-00:10:00&order=column")
         res = json.loads(res.data.decode())
         assert res["code"] == 200
@@ -159,10 +189,7 @@ class TestDataAPI:
         assert len(res["results"]["data"]["air_temp"]) == 9
         assert len(res["results"]["timestamps"]) == 9
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_wind_speed_direction_json(self, query_func, client):
-        r = fake_data("1m", {("aoss", "tower"): ["time", "wind_speed", "wind_direction", "wind_east", "wind_north"]}, 9)
-        query_func.return_value = r
+    def test_wind_speed_direction_json(self, client, influxdb_wind_fields_9_values):
         res = client.get(
             "/api/data.json?symbols=aoss.tower.wind_speed:aoss.tower.wind_direction&begin=-00:10:00&order=column"
         )
@@ -173,9 +200,8 @@ class TestDataAPI:
         assert "aoss.tower.wind_speed" in res["results"]["data"]
         assert len(list(res["results"]["data"].keys())) == 2
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_one_symbol_two_insts_json_row(self, query_func, client):
-        r = fake_data(
+    def test_one_symbol_two_insts_json_row(self, client, mock_influxdb_query):
+        fake_result = _fake_data(
             "1m",
             {
                 ("aoss", "tower"): ["time", "air_temp"],
@@ -183,9 +209,9 @@ class TestDataAPI:
             },
             9,
         )
-        query_func.return_value = r
-        # row should be the default
-        res = client.get("/api/data.json?symbols=aoss.tower.air_temp:mendota.buoy.air_temp&begin=-00:10:00")
+        with mock_influxdb_query(fake_result):
+            # row should be the default
+            res = client.get("/api/data.json?symbols=aoss.tower.air_temp:mendota.buoy.air_temp&begin=-00:10:00")
         res = json.loads(res.data.decode())
         assert res["code"] == 200
         assert res["num_results"] == 9
@@ -194,9 +220,8 @@ class TestDataAPI:
         assert len(res["results"]["data"]) == 9
         assert len(res["results"]["data"][0]) == 2
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_one_symbol_three_insts_json_row(self, query_func, client):
-        r = fake_data(
+    def test_one_symbol_three_insts_json_row(self, client, mock_influxdb_query):
+        fake_result = _fake_data(
             "1m",
             {
                 ("site1", "inst1"): ["time", "air_temp"],
@@ -205,7 +230,6 @@ class TestDataAPI:
             },
             9,
         )
-        query_func.return_value = r
         # row should be the default
         from metobsapi.util.data_responses import SYMBOL_TRANSLATIONS as st
 
@@ -213,7 +237,7 @@ class TestDataAPI:
         st[("site1", "inst1")] = st[("aoss", "tower")]
         st[("site2", "inst2")] = st[("aoss", "tower")]
         st[("site3", "inst3")] = st[("aoss", "tower")]
-        with mock.patch("metobsapi.util.data_responses.SYMBOL_TRANSLATIONS", st):
+        with mock.patch("metobsapi.util.data_responses.SYMBOL_TRANSLATIONS", st), mock_influxdb_query(fake_result):
             res = client.get(
                 "/api/data.json?symbols=site1.inst1.air_temp:site2.inst2.air_temp:site3.inst3.air_temp&begin=-00:10:00"
             )
@@ -225,10 +249,7 @@ class TestDataAPI:
             assert len(res["results"]["data"]) == 9
             assert len(res["results"]["data"][0]) == 3
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_one_symbol_csv(self, query_func, client):
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp"]}, 9)
-        query_func.return_value = r
+    def test_one_symbol_csv(self, client, influxdb_air_temp_9_values):
         # row should be the default
         res = client.get("/api/data.csv?symbols=aoss.tower.air_temp&begin=-00:10:00")
         res = res.data.decode()
@@ -239,12 +260,9 @@ class TestDataAPI:
         assert len(lines[5].split(",")) == 2
         assert "# code: 200" in res
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_one_symbol_xml(self, query_func, client):
+    def test_one_symbol_xml(self, client, influxdb_air_temp_9_values):
         from xml.dom.minidom import parseString
 
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp"]}, 9)
-        query_func.return_value = r
         # row should be the default
         res = client.get("/api/data.xml?symbols=aoss.tower.air_temp&begin=-00:10:00")
         res = parseString(res.data.decode())
@@ -253,11 +271,8 @@ class TestDataAPI:
         # data rows
         assert len(res.childNodes[0].childNodes[1].childNodes) == 9
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_three_symbol_csv(self, query_func, client):
+    def test_three_symbol_csv(self, client, influxdb_3_symbols_9_values):
         """Test that multiple channels in a CSV file are structured properly."""
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp", "rel_hum", "wind_speed"]}, 9)
-        query_func.return_value = r
         # row should be the default
         res = client.get(
             "/api/data.csv?symbols=aoss.tower.air_temp:" "aoss.tower.rel_hum:aoss.tower.wind_speed&begin=-00:10:00"
@@ -270,11 +285,8 @@ class TestDataAPI:
         assert len(lines[5].split(",")) == 4
         assert "# code: 200" in res
 
-    @mock.patch("metobsapi.data_api.query")
-    def test_three_symbol_csv_repeat(self, query_func, client):
+    def test_three_symbol_csv_repeat(self, client, influxdb_3_symbols_9_values):
         """Test that multiple channels in a CSV file are structured properly."""
-        r = fake_data("1m", {("aoss", "tower"): ["time", "air_temp", "rel_hum", "wind_speed"]}, 9)
-        query_func.return_value = r
         # row should be the default
         res = client.get(
             "/api/data.csv?symbols=aoss.tower.air_temp:" "aoss.tower.air_temp:aoss.tower.air_temp&begin=-00:10:00"
@@ -291,7 +303,7 @@ class TestDataAPI:
     # @mock.patch('metobsapi.data_api.query')
     # def test_jsonp_bad_symbol_400(self, query_func, client):
     #     XXX: Not currently possible with flask-json
-    #     r = fake_data('1m', {('aoss', 'tower'): ['time', 'air_temp']}, 9)
+    #     r = _fake_data('1m', {('aoss', 'tower'): ['time', 'air_temp']}, 9)
     #     query_func.return_value = r
     #     # row should be the default
     #     res = client.get('/api/data.json?site=aoss&inst=tower&symbols=bad&begin=-00:10:00&callback=test')