From 6c15ee8517d33741f4e930d56736af5c3c2dcef7 Mon Sep 17 00:00:00 2001
From: lilydjwg <lilydjwg@gmail.com>
Date: Tue, 8 May 2018 18:34:32 +0800
Subject: [PATCH] expose HTTPError, and raise it if status >= 400 for aiohttp

---
 nvchecker/source/aiohttp_httpclient.py | 13 +++++++++++--
 nvchecker/source/tornado_httpclient.py |  3 ++-
 tests/test_keyfile.py                  |  4 +++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/nvchecker/source/aiohttp_httpclient.py b/nvchecker/source/aiohttp_httpclient.py
index 9bd9e59..85cff0a 100644
--- a/nvchecker/source/aiohttp_httpclient.py
+++ b/nvchecker/source/aiohttp_httpclient.py
@@ -5,14 +5,23 @@ import atexit
 import aiohttp
 connector = aiohttp.TCPConnector(limit=20)
 
-__all__ = ['session']
+__all__ = ['session', 'HTTPError']
+
+class HTTPError(Exception):
+    def __init__(self, code, message):
+        self.code = code
+        self.message = message
 
 class BetterClientSession(aiohttp.ClientSession):
     async def _request(self, *args, **kwargs):
         if hasattr(self, "nv_config") and self.nv_config.get("proxy"):
             kwargs.setdefault("proxy", self.nv_config.get("proxy"))
 
-        return await super(BetterClientSession, self)._request(*args, **kwargs)
+        res = await super(BetterClientSession, self)._request(
+            *args, **kwargs)
+        if res.status >= 400:
+            raise HTTPError(res.status, res.reason)
+        return res
 
 session = BetterClientSession(connector=connector, read_timeout=10, conn_timeout=5)
 atexit.register(session.close)
diff --git a/nvchecker/source/tornado_httpclient.py b/nvchecker/source/tornado_httpclient.py
index a87d098..53b46e4 100644
--- a/nvchecker/source/tornado_httpclient.py
+++ b/nvchecker/source/tornado_httpclient.py
@@ -5,6 +5,7 @@ import json
 from urllib.parse import urlencode
 
 from tornado.httpclient import AsyncHTTPClient, HTTPRequest, HTTPResponse
+from tornado.httpclient import HTTPError
 from tornado.platform.asyncio import AsyncIOMainLoop, to_asyncio_future
 AsyncIOMainLoop().install()
 
@@ -14,7 +15,7 @@ try:
 except ImportError:
   pycurl = None
 
-__all__ = ['session']
+__all__ = ['session', 'HTTPError']
 
 client = AsyncHTTPClient()
 HTTP2_AVAILABLE = None if pycurl else False
diff --git a/tests/test_keyfile.py b/tests/test_keyfile.py
index 3e51dea..c237b57 100644
--- a/tests/test_keyfile.py
+++ b/tests/test_keyfile.py
@@ -3,6 +3,8 @@
 
 import tempfile
 
+from nvchecker.source import HTTPError
+
 import pytest
 pytestmark = [pytest.mark.asyncio]
 
@@ -31,5 +33,5 @@ keyfile = {f.name}
 
     try:
       await run_source(test_conf)
-    except Exception as e:
+    except HTTPError as e:
       assert e.code == 401