Move to aiohttp for GeoIP requests

GeoIP requests used to run on the default executor thread and would
prevent the application from exiting if the GeoIP service would not
respond quickly enough. We witnessed an obvious impact during an
incident on geoip.ubuntu.com.

Move to aiohttp so that the HTTP calls are non blocking.

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

Makefile

@@ -31,7 +31,7 @@ aptdeps:
 		python3-requests-unixsocket python3-jsonschema python3-apport \
 		python3-bson xorriso isolinux python3-aiohttp cloud-init ssh-import-id \
 		curl jq build-essential python3-pytest python3-async-timeout \
-	        language-selector-common fuseiso python3-pytest-xdist
+	        language-selector-common fuseiso python3-pytest-xdist python3-aioresponses
 
 .PHONY: install_deps
 install_deps: aptdeps gitdeps

subiquity/server/geoip.py

@@ -14,13 +14,12 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from abc import ABC, abstractmethod
+import aiohttp
 import logging
 import enum
-import requests
 from xml.etree import ElementTree
 
 from subiquitycore.async_helpers import (
-    run_in_thread,
     SingleInstanceTask,
 )
 
@@ -77,12 +76,12 @@ class HTTPGeoIPStrategy(GeoIPStrategy):
     geoip.ubuntu.com service. """
     async def get_response(self) -> str:
         try:
-            response = await run_in_thread(
-                requests.get, "https://geoip.ubuntu.com/lookup")
-            response.raise_for_status()
-        except requests.exceptions.RequestException as e:
+            async with aiohttp.ClientSession() as session:
+                async with session.get("https://geoip.ubuntu.com/lookup") as response:
+                    response.raise_for_status()
+                    return await response.text()
+        except aiohttp.ClientError as e:
             raise LookupError from e
-        return response.text
 
 
 class GeoIP:

subiquity/server/tests/test_geoip.py

@@ -49,30 +49,17 @@ empty_tz = '<Response><TimeZone></TimeZone></Response>'
 empty_cc = '<Response><CountryCode></CountryCode></Response>'
 
 
-class MockGeoIPResponse:
-    def __init__(self, text, status_code=200):
-        self.text = text
-        self.status_code = status_code
-
-    def raise_for_status(self, *args, **kwargs):
-        pass
-
-
-def requests_get_factory(text):
-    def requests_get(*args, **kwargs):
-        return MockGeoIPResponse(text)
-    return requests_get
-
-
 class TestGeoIP(SubiTestCase):
-    @mock.patch('requests.get', new=requests_get_factory(xml))
     def setUp(self):
         strategy = HTTPGeoIPStrategy()
         self.geoip = GeoIP(make_app(), strategy)
 
         async def fn():
             self.assertTrue(await self.geoip.lookup())
-        run_coro(fn())
+
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=xml)
+            run_coro(fn())
 
     def test_countrycode(self):
         self.assertEqual("us", self.geoip.countrycode)
@@ -86,37 +73,42 @@ class TestGeoIPBadData(SubiTestCase):
         strategy = HTTPGeoIPStrategy()
         self.geoip = GeoIP(make_app(), strategy)
 
-    @mock.patch('requests.get', new=requests_get_factory(partial))
     def test_partial_reponse(self):
         async def fn():
             self.assertFalse(await self.geoip.lookup())
-        run_coro(fn())
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=partial)
+            run_coro(fn())
 
-    @mock.patch('requests.get', new=requests_get_factory(incomplete))
     def test_incomplete(self):
         async def fn():
             self.assertFalse(await self.geoip.lookup())
-        run_coro(fn())
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=incomplete)
+            run_coro(fn())
         self.assertIsNone(self.geoip.countrycode)
         self.assertIsNone(self.geoip.timezone)
 
-    @mock.patch('requests.get', new=requests_get_factory(long_cc))
     def test_long_cc(self):
         async def fn():
             self.assertFalse(await self.geoip.lookup())
-        run_coro(fn())
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=long_cc)
+            run_coro(fn())
         self.assertIsNone(self.geoip.countrycode)
 
-    @mock.patch('requests.get', new=requests_get_factory(empty_cc))
     def test_empty_cc(self):
         async def fn():
             self.assertFalse(await self.geoip.lookup())
-        run_coro(fn())
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=empty_cc)
+            run_coro(fn())
         self.assertIsNone(self.geoip.countrycode)
 
-    @mock.patch('requests.get', new=requests_get_factory(empty_tz))
     def test_empty_tz(self):
         async def fn():
             self.assertFalse(await self.geoip.lookup())
-        run_coro(fn())
+        with aioresponses() as mocked:
+            mocked.get("https://geoip.ubuntu.com/lookup", body=empty_tz)
+            run_coro(fn())
         self.assertIsNone(self.geoip.timezone)