Skip to content

Commit

Permalink
Refactor: Updated the code & test for FTP_TLS
Browse files Browse the repository at this point in the history
  • Loading branch information
@fazledyn-or
fazledyn-or committed Oct 9, 2023
1 parent 84a00d8 commit 099196a
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 170 deletions.
16 changes: 10 additions & 6 deletions storages/backends/ftp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,7 @@
# In models.py you can write:
# from FTPStorage import FTPStorage
# fs = FTPStorage()
# For a TLS configuration, you must use it below:
# fs = FTPStorage(secure=True)
# For a TLS configuration, you must use 'ftps' protocol
# class FTPTest(models.Model):
# file = models.FileField(upload_to='a/b/c/', storage=fs)

Expand All @@ -39,7 +38,7 @@ class FTPStorageException(Exception):
class FTPStorage(Storage):
"""FTP Storage class for Django pluggable storage system."""

def __init__(self, location=None, base_url=None, encoding=None, secure=False):
def __init__(self, location=None, base_url=None, encoding=None):
location = location or setting("FTP_STORAGE_LOCATION")
if location is None:
raise ImproperlyConfigured(
Expand All @@ -53,7 +52,6 @@ def __init__(self, location=None, base_url=None, encoding=None, secure=False):
self._config = self._decode_location(location)
self._base_url = base_url
self._connection = None
self._secure = secure

def _decode_location(self, location):
"""Return splitted configuration data from location."""
Expand All @@ -69,6 +67,12 @@ def _decode_location(self, location):
config["active"] = True
else:
config["active"] = False

if splitted_url.scheme == "ftps":
config["secure"] = True
else:
config["secure"] = False

config["path"] = splitted_url.path
config["host"] = splitted_url.hostname
config["user"] = splitted_url.username
Expand All @@ -87,12 +91,12 @@ def _start_connection(self):

# Real reconnect
if self._connection is None:
ftp = ftplib.FTP_TLS() if self._secure else ftplib.FTP()
ftp = ftplib.FTP_TLS() if self._config["secure"] else ftplib.FTP()
ftp.encoding = self.encoding
try:
ftp.connect(self._config["host"], self._config["port"])
ftp.login(self._config["user"], self._config["passwd"])
if self._secure:
if self._config["secure"]:
ftp.prot_p()
if self._config["active"]:
ftp.set_pasv(False)
Expand Down
171 changes: 7 additions & 164 deletions tests/test_ftp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ def test_decode_location(self):
"active": False,
"path": "/",
"port": 2121,
"secure": False
}
self.assertEqual(config, wanted_config)
# Test active FTP
Expand All @@ -62,6 +63,7 @@ def test_decode_location(self):
"active": True,
"path": "/",
"port": 2121,
"secure": False
}
self.assertEqual(config, wanted_config)

Expand Down Expand Up @@ -246,15 +248,15 @@ def test_close(self, mock_ftp, mock_storage):

class FTPTLSTest(TestCase):
def setUp(self):
self.storage = ftp.FTPStorage(location=URL_TLS, secure=True)
self.storage = ftp.FTPStorage(location=URL_TLS)

def test_init_no_location(self):
with self.assertRaises(ImproperlyConfigured):
ftp.FTPStorage(secure=True)
ftp.FTPStorage()

@patch("storages.backends.ftp.setting", return_value=URL_TLS)
def test_init_location_from_setting(self, mock_setting):
storage = ftp.FTPStorage(secure=True)
storage = ftp.FTPStorage()
self.assertTrue(mock_setting.called)
self.assertEqual(storage.location, URL_TLS)

Expand All @@ -267,6 +269,7 @@ def test_decode_location(self):
"active": False,
"path": "/",
"port": 2121,
"secure": True
}
self.assertEqual(config, wanted_config)

Expand All @@ -284,166 +287,6 @@ def test_start_connection(self, mock_ftp):
self.storage._start_connection()
self.assertIsNotNone(self.storage._connection)
# Start active
storage = ftp.FTPStorage(location=URL_TLS, secure=True)
storage = ftp.FTPStorage(location=URL_TLS)
storage._start_connection()

@patch("ftplib.FTP_TLS", **{"return_value.pwd.side_effect": IOError()})
def test_start_connection_timeout(self, mock_ftp):
self.storage._start_connection()
self.assertIsNotNone(self.storage._connection)

@patch("ftplib.FTP_TLS", **{"return_value.connect.side_effect": IOError()})
def test_start_connection_error(self, mock_ftp):
with self.assertRaises(ftp.FTPStorageException):
self.storage._start_connection()

@patch("ftplib.FTP_TLS", **{"return_value.quit.return_value": None})
def test_disconnect(self, mock_ftp_quit):
self.storage._start_connection()
self.storage.disconnect()
self.assertIsNone(self.storage._connection)

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
def test_mkremdirs(self, mock_ftp):
self.storage._start_connection()
self.storage._mkremdirs("foo/bar")

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
def test_mkremdirs_n_subdirectories(self, mock_ftp):
self.storage._start_connection()
self.storage._mkremdirs("foo/bar/null")

@patch(
"ftplib.FTP_TLS",
**{
"return_value.pwd.return_value": "foo",
"return_value.storbinary.return_value": None,
},
)
def test_put_file(self, mock_ftp):
self.storage._start_connection()
self.storage._put_file("foo", File(io.BytesIO(b"foo"), "foo"))

@patch(
"ftplib.FTP_TLS",
**{
"return_value.pwd.return_value": "foo",
"return_value.storbinary.side_effect": IOError(),
},
)
def test_put_file_error(self, mock_ftp):
self.storage._start_connection()
with self.assertRaises(ftp.FTPStorageException):
self.storage._put_file("foo", File(io.BytesIO(b"foo"), "foo"))

def test_open(self):
remote_file = self.storage._open("foo")
self.assertIsInstance(remote_file, ftp.FTPStorageFile)

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
def test_read(self, mock_ftp):
self.storage._start_connection()
self.storage._read("foo")

@patch("ftplib.FTP_TLS", **{"return_value.pwd.side_effect": IOError()})
def test_read2(self, mock_ftp):
self.storage._start_connection()
with self.assertRaises(ftp.FTPStorageException):
self.storage._read("foo")

@patch(
"ftplib.FTP_TLS",
**{
"return_value.pwd.return_value": "foo",
"return_value.storbinary.return_value": None,
},
)
def test_save(self, mock_ftp):
self.storage._save("foo", File(io.BytesIO(b"foo"), "foo"))

@patch("ftplib.FTP_TLS", **{"return_value.retrlines": list_retrlines})
def test_listdir(self, mock_retrlines):
dirs, files = self.storage.listdir("/")
self.assertEqual(len(dirs), 1)
self.assertEqual(dirs, ["dir"])
self.assertEqual(len(files), 2)
self.assertEqual(sorted(files), sorted(["fi", "fi2"]))

@patch("ftplib.FTP_TLS", **{"return_value.retrlines.side_effect": IOError()})
def test_listdir_error(self, mock_ftp):
with self.assertRaises(ftp.FTPStorageException):
self.storage.listdir("/")

@patch("ftplib.FTP_TLS", **{"return_value.nlst.return_value": ["foo", "foo2"]})
def test_exists(self, mock_ftp):
self.assertTrue(self.storage.exists("foo"))
self.assertFalse(self.storage.exists("bar"))

@patch("ftplib.FTP_TLS", **{"return_value.nlst.side_effect": IOError()})
def test_exists_error(self, mock_ftp):
with self.assertRaises(ftp.FTPStorageException):
self.storage.exists("foo")

@patch(
"ftplib.FTP_TLS",
**{
"return_value.delete.return_value": None,
"return_value.nlst.return_value": ["foo", "foo2"],
},
)
def test_delete(self, mock_ftp):
self.storage.delete("foo")
self.assertTrue(mock_ftp.return_value.delete.called)

@patch("ftplib.FTP_TLS", **{"return_value.retrlines": list_retrlines})
def test_size(self, mock_ftp):
self.assertEqual(1024, self.storage.size("fi"))
self.assertEqual(2048, self.storage.size("fi2"))
self.assertEqual(0, self.storage.size("bar"))

@patch("ftplib.FTP_TLS", **{"return_value.retrlines.side_effect": IOError()})
def test_size_error(self, mock_ftp):
self.assertEqual(0, self.storage.size("foo"))

def test_url(self):
with self.assertRaises(ValueError):
self.storage._base_url = None
self.storage.url("foo")
self.storage = ftp.FTPStorage(location=URL_TLS, base_url="http://foo.bar/", secure=True)
self.assertEqual("http://foo.bar/foo", self.storage.url("foo"))


class FTPTLSStorageFileTest(TestCase):
def setUp(self):
self.storage = ftp.FTPStorage(location=URL_TLS, secure=True)

@patch("ftplib.FTP_TLS", **{"return_value.retrlines": list_retrlines})
def test_size(self, mock_ftp):
file_ = ftp.FTPStorageFile("fi", self.storage, "wb")
self.assertEqual(file_.size, 1024)

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
@patch("storages.backends.ftp.FTPStorage._read", return_value=io.BytesIO(b"foo"))
def test_readlines(self, mock_ftp, mock_storage):
file_ = ftp.FTPStorageFile("fi", self.storage, "wb")
self.assertEqual([b"foo"], file_.readlines())

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
@patch("storages.backends.ftp.FTPStorage._read", return_value=io.BytesIO(b"foo"))
def test_read(self, mock_ftp, mock_storage):
file_ = ftp.FTPStorageFile("fi", self.storage, "wb")
self.assertEqual(b"foo", file_.read())

def test_write(self):
file_ = ftp.FTPStorageFile("fi", self.storage, "wb")
file_.write(b"foo")
file_.seek(0)
self.assertEqual(file_.file.read(), b"foo")

@patch("ftplib.FTP_TLS", **{"return_value.pwd.return_value": "foo"})
@patch("storages.backends.ftp.FTPStorage._read", return_value=io.BytesIO(b"foo"))
def test_close(self, mock_ftp, mock_storage):
file_ = ftp.FTPStorageFile("fi", self.storage, "wb")
file_.is_dirty = True
file_.read()
file_.close()

0 comments on commit 099196a

Please sign in to comment.