def test_ftp(mocker):
ftp = mocker.patch("calculadora_do_cidadao.download.FTP")
path = mocker.MagicMock()
download = Download("ftp://here.comes/a/fancy/url")
download.ftp(path)
ftp.assert_called_once_with("here.comes")
ftp.return_value.__enter__.return_value.login.assert_called_once_with()
ftp.return_value.__enter__.return_value.retrbinary.assert_called_once_with(
"RETR /a/fancy/url",
path.open.return_value.__enter__.return_value.write)
def test_http_post(mocker):
session = mocker.patch("calculadora_do_cidadao.download.Session")
session.return_value.post.return_value.content = b"42"
path = mocker.MagicMock()
download = Download("https://here.comes/a/fancy.url", post_data={"test": 42})
download.http(path)
session.assert_called_once_with()
session.return_value.post.assert_called_once_with(
"https://here.comes/a/fancy.url", data={"test": 42}
)
path.write_bytes.assert_called_once_with(b"42")
def test_http_get(mocker):
session = mocker.patch("calculadora_do_cidadao.download.Session")
session.return_value.get.return_value.content = b"42"
jar = mocker.patch("calculadora_do_cidadao.download.cookiejar_from_dict")
jar.return_value = "my-cookie-jar"
path = mocker.MagicMock()
download = Download("https://here.comes/a/fancy.url", cookies={"test": 42})
download.http(path)
jar.assert_called_once_with({"test": 42})
assert session.return_value.cookies == "my-cookie-jar"
session.assert_called_once_with()
session.return_value.get.assert_called_once_with("https://here.comes/a/fancy.url")
path.write_bytes.assert_called_once_with(b"42")
def download(self) -> IndexesGenerator:
"""Wrapper to use the `Download` class and pipe the result to `rows`
imported method, yielding a series of rows parsed by `rows`."""
download = Download(self.url, self.should_unzip, self.cookies)
with download() as path:
for kwargs in self.import_kwargs:
for data in self.read_from(path, **kwargs):
yield from (row for row in self.serialize(data) if row)
def test_post_procesing(mocker, broken_table):
mocker.patch.object(Download, "http", return_value=broken_table)
download = Download(
"http://here.comes/a/fancy/url.zip", post_processing=lambda b: b"<table>" + b
)
with download() as path:
download.http.assert_called_once()
assert path.read_text().startswith("<table>")
assert path.read_text().endswith("</table>\n")
def download(self) -> Iterator[Tuple[date, Decimal]]:
download = Download(self.url)
with download() as path:
table = import_from_xls(path, **self.import_kwargs)
rows = (self.serialize(row) for row in table)
yield from (row for row in rows if row)
def test_unzip(zip_file):
assert Download.unzip(zip_file).read_bytes() == b"42"
def test_download_not_implemented():
expected = r"No method implemented for tcp\." # this is a regex
with raises(DownloadMethodNotImplementedError, match=expected):
Download("tcp://here.comes/a/fancy/url.zip")
def test_download(zip_file, mocker):
mocker.patch.object(Download, "ftp", return_value=zip_file)
download = Download("ftp://here.comes/a/fancy/url.zip", should_unzip=True)
with download() as path:
download.ftp.assert_called_once()
assert path.read_bytes() == b"42"
def download(self) -> IndexesGenerator:
download = Download(self.url, self.should_unzip, self.cookies)
with download() as path:
for kwargs in self.import_kwargs:
for data in self.read_from(path, **kwargs):
yield from (row for row in self.serialize(data) if row)