diff --git a/apt-deps.txt b/apt-deps.txt
index eab49e9b..03e127f1 100644
--- a/apt-deps.txt
+++ b/apt-deps.txt
@@ -15,6 +15,7 @@ os-prober
pep8
pkg-config
python3-aiohttp
+python3-aioresponses
python3-apport
python3-async-timeout
python3-attr
diff --git a/subiquity/server/geoip.py b/subiquity/server/geoip.py
index 0f9eeef0..1d38fc68 100644
--- a/subiquity/server/geoip.py
+++ b/subiquity/server/geoip.py
@@ -14,13 +14,12 @@
# along with this program. If not, see .
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,
)
@@ -76,13 +75,14 @@ class HTTPGeoIPStrategy(GeoIPStrategy):
""" HTTP implementation to retrieve GeoIP information. We use the
geoip.ubuntu.com service. """
async def get_response(self) -> str:
+ url = "https://geoip.ubuntu.com/lookup"
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(url) as response:
+ response.raise_for_status()
+ return await response.text()
+ except aiohttp.ClientError as e:
raise LookupError from e
- return response.text
class GeoIP:
diff --git a/subiquity/server/tests/test_geoip.py b/subiquity/server/tests/test_geoip.py
index ae15e5f8..835292e8 100644
--- a/subiquity/server/tests/test_geoip.py
+++ b/subiquity/server/tests/test_geoip.py
@@ -13,7 +13,7 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
-import mock
+from aioresponses import aioresponses
from subiquitycore.tests import SubiTestCase
from subiquitycore.tests.mocks import make_app
@@ -47,30 +47,17 @@ empty_tz = ''
empty_cc = ''
-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)
@@ -84,37 +71,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)