From 10af626767e124ede1d809eb4b7216956a98bd91 Mon Sep 17 00:00:00 2001 From: Olivier Gayot Date: Wed, 18 May 2022 16:26:30 +0200 Subject: [PATCH] geoip: move from requests to aiohttp again python3-aiohttp has been fixed and released in focal-updates. New Subiquity builds should not be affected by LP: #1969393 anymore. This reverts commit 413cc873379fa15964575047e3c76f99439a8e54. Signed-off-by: Olivier Gayot --- apt-deps.txt | 1 + subiquity/server/geoip.py | 14 ++++---- subiquity/server/tests/test_geoip.py | 48 ++++++++++++---------------- 3 files changed, 28 insertions(+), 35 deletions(-) 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)