Add error handling to 'download_file' (#64)

Author: ostefano

tests/test_client.py

@@ -20,6 +20,7 @@
 
 import pytest
 
+from vt import APIError
 from vt import Client
 from vt import FeedType
 from vt import Object
@@ -279,6 +280,31 @@ def test_download_file(httpserver):
       assert f.read() == b'filecontent'
 
 
+def test_download_file_with_error(httpserver):
+
+  httpserver.expect_request(
+    '/api/v3/files/01020304050607080900a0b0c0ddead/download',
+    method='GET',
+    headers={'X-Apikey': 'dummy_api_key'}
+  ).respond_with_data(
+    status=404,
+    content_type="application/json",
+    response_data=json.dumps({
+      "error": {
+        "code": "NotFoundError",
+        "message": "Resource not found."
+      }
+    })
+  )
+
+  with pytest.raises(APIError) as e_info:
+    with new_client(httpserver) as client:
+      with io.BytesIO() as f:
+        client.download_file('01020304050607080900a0b0c0ddead', f)
+  assert e_info.value.args[0] == "NotFoundError"
+  assert e_info.value.args[1] == "Resource not found."
+
+
 def test_scan_file(httpserver):
 
   upload_url = (

vt/client.py

@@ -308,6 +308,9 @@ def download_file(self, hash, file):
   async def download_file_async(self, hash, file):
     """Like :func:`download_file` but returns a coroutine."""
     response = await self.get_async(f'/files/{hash}/download')
+    error = await self.get_error_async(response)
+    if error:
+      raise error
     while True:
       chunk = await response.content.read_async(1024*1024)
       if not chunk:

vt/version.py

@@ -1 +1 @@
-__version__ = '0.7.0'
+__version__ = '0.7.1'