From cba93954cd09bd11fcd99466717d340b23c94a10 Mon Sep 17 00:00:00 2001 From: Jimmy Date: Tue, 2 Jan 2018 15:56:28 +1300 Subject: [PATCH] Allow download_stub test fixture to handle file targets. --- tests/unit/browser/test_adblock.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/unit/browser/test_adblock.py b/tests/unit/browser/test_adblock.py index 09161e806..a6bbcd8ce 100644 --- a/tests/unit/browser/test_adblock.py +++ b/tests/unit/browser/test_adblock.py @@ -23,12 +23,13 @@ import os.path import zipfile import shutil import logging +import contextlib import pytest from PyQt5.QtCore import pyqtSignal, QUrl, QObject -from qutebrowser.browser import adblock +from qutebrowser.browser import adblock, downloads from qutebrowser.utils import objreg pytestmark = pytest.mark.usefixtures('qapp', 'config_tmpdir') @@ -89,14 +90,27 @@ class FakeDownloadManager: def __init__(self, tmpdir): self._tmpdir = tmpdir + @contextlib.contextmanager + def _open_fileobj(self, target): + """Ensure a DownloadTarget's fileobj attribute is available.""" + if isinstance(target, downloads.FileDownloadTarget): + target.fileobj = open(target.filename, 'wb') + try: + yield target.fileobj + finally: + target.fileobj.close() + else: + yield target.fileobj + def get(self, url, target, **kwargs): """Return a FakeDownloadItem instance with a fileobj. The content is copied from the file the given url links to. """ - download_item = FakeDownloadItem(target.fileobj, name=url.path()) - with (self._tmpdir / url.path()).open('rb') as fake_url_file: - shutil.copyfileobj(fake_url_file, download_item.fileobj) + with self._open_fileobj(target): + download_item = FakeDownloadItem(target.fileobj, name=url.path()) + with (self._tmpdir / url.path()).open('rb') as fake_url_file: + shutil.copyfileobj(fake_url_file, download_item.fileobj) return download_item