Aioftp (#398)

* coverted from from ftp to aioftp

* required changes made

* async function call + tests

* test + ftp call fix

* aioftp exception fixed

* test_fix_3

* minor corrections

* updated tests and rfi

* switch to asynctest

* Fix pep8 issues

* Mock Async FTP connection

Co-authored-by: steve7158 <stevejmotha@gmail.com>

requirements.txt

@@ -20,3 +20,5 @@ mako
 pyjwt
 pyyaml
 sqlalchemy
+aioftp
+asynctest

tanner/emulators/rfi.py

@@ -1,13 +1,12 @@
 import asyncio
-import ftplib
+import aioftp
 import hashlib
 import logging
 import os
 import re
 import ssl
 import time
-from concurrent.futures import ThreadPoolExecutor
-
+from aioftp.errors import AIOFTPException, StatusCodeError, PathIsNotAbsolute, PathIOError, NoAvailablePort
 import aiohttp
 import yarl
 
@@ -26,7 +25,6 @@ def __init__(self, root_dir, loop=None, allow_insecure=False):
     async def download_file(self, path):
         file_name = None
         url = re.match(patterns.REMOTE_FILE_URL, path)
-
         if url is None:
             return None
         url = url.group(1)
@@ -36,9 +34,7 @@ async def download_file(self, path):
             os.makedirs(self.script_dir)
 
         if url.scheme == "ftp":
-            pool = ThreadPoolExecutor()
-            ftp_future = self._loop.run_in_executor(pool, self.download_file_ftp, url)
-            file_name = await ftp_future
+            file_name = await self.download_file_ftp(url)
 
         else:
             ssl_context = False if self.allow_insecure else ssl.create_default_context()
@@ -56,20 +52,17 @@ async def download_file(self, path):
                     rfile.write(data.encode('utf-8'))
         return file_name
 
-    def download_file_ftp(self, url):
+    async def download_file_ftp(self, url):
         host = url.host
-        ftp_path = url.path.rsplit('/', 1)[0][1:]
+        ftp_path = url.path
         name = url.name
         try:
-            ftp = ftplib.FTP(host)
-            ftp.login()
-            ftp.cwd(ftp_path)
-            tmp_filename = name + str(time.time())
-            file_name = hashlib.md5(tmp_filename.encode('utf-8')).hexdigest()
-            with open(os.path.join(self.script_dir, file_name), 'wb') as ftp_script:
+            async with aioftp.ClientSession(host) as ftp:
+                tmp_filename = name + str(time.time())
+                file_name = hashlib.md5(tmp_filename.encode('utf-8')).hexdigest()
                 self.logger.debug('Saving the FTP file as %s', os.path.join(self.script_dir, file_name))
-                ftp.retrbinary('RETR %s' % name, ftp_script.write)
-        except ftplib.all_errors as ftp_errors:
+                await ftp.download(ftp_path, os.path.join(self.script_dir, file_name), write_into=True)
+        except (AIOFTPException, StatusCodeError, PathIsNotAbsolute, PathIOError, NoAvailablePort) as ftp_errors:
             self.logger.exception("Problem with ftp download %s", ftp_errors)
             return None
         else:

tanner/tests/test_rfi_emulation.py

@@ -3,42 +3,44 @@
 from unittest import mock
 from tanner.emulators import rfi
 import yarl
+from tanner.utils.asyncmock import AysncFTPMock
+from asynctest import CoroutineMock
 
 
 class TestRfiEmulator(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(None)
-        self.handler = rfi.RfiEmulator('/tmp/', loop=self.loop)
+        self.handler = rfi.RfiEmulator("/tmp/", loop=self.loop)
 
     def test_http_download(self):
-        path = 'http://example.com'
+        path = "http://example.com"
         data = self.loop.run_until_complete(self.handler.download_file(path))
         self.assertIsNotNone(data)
 
     def test_http_download_fail(self):
-        path = 'http://foobarfvfd'
+        path = "http://foobarfvfd"
         filename = self.loop.run_until_complete(self.handler.download_file(path))
         self.assertIsNone(filename)
 
     def test_ftp_download(self):
-        self.handler.download_file_ftp = mock.MagicMock()
-        path = 'ftp://mirror.yandex.ru/archlinux/lastupdate'
+        self.handler.download_file_ftp = CoroutineMock()
+        path = "ftp://mirror.yandex.ru/archlinux/lastupdate"
         data = self.loop.run_until_complete(self.handler.download_file(path))
-        self.handler.download_file_ftp.assert_called_with(yarl.URL(path))
-
-    def test_ftp_download_fail(self):
-        path = 'ftp://mirror.yandex.ru/archlinux/foobar'
+        self.handler.download_file_ftp.assert_awaited_with(yarl.URL(path))
 
+    async def test_ftp_download_fail(self):
+        self.handler.download_file = AysncFTPMock()
+        path = "ftp://mirror.yandex.ru/archlinux/foobar"
         with self.assertLogs():
-            self.loop.run_until_complete(self.handler.download_file(path))
+            await self.handler.download_file(path)
 
     def test_get_result_fail(self):
         data = "test data"
         result = self.loop.run_until_complete(self.handler.get_rfi_result(data))
         self.assertIsNone(result)
 
     def test_invalid_scheme(self):
-        path = 'file://mirror.yandex.ru/archlinux/foobar'
+        path = "file://mirror.yandex.ru/archlinux/foobar"
         data = self.loop.run_until_complete(self.handler.download_file(path))
         self.assertIsNone(data)

tanner/utils/asyncmock.py

@@ -1,5 +1,6 @@
 import asyncio
 from unittest.mock import Mock
+from aioftp.errors import AIOFTPException
 
 
 class AsyncMock(Mock):  # custom function defined to mock asyncio coroutines
@@ -13,3 +14,9 @@ async def coro():
 
     def __await__(self):
         return self().__await__()
+
+
+class AysncFTPMock(Mock):
+
+    def __aenter__():
+        raise AIOFTPException