diff options
Diffstat (limited to 'test_upscale_qt.py')
| -rw-r--r-- | test_upscale_qt.py | 84 |
1 files changed, 84 insertions, 0 deletions
diff --git a/test_upscale_qt.py b/test_upscale_qt.py new file mode 100644 index 0000000..526dc37 --- /dev/null +++ b/test_upscale_qt.py @@ -0,0 +1,84 @@ +# upscale-qt - tests for the non-GUI helpers and a GUI smoke test +# Copyright (C) 2026 Danilo M. <danix@danix.xyz> +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 2 as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + +import os +import tempfile +from importlib.machinery import SourceFileLoader +from importlib.util import module_from_spec, spec_from_loader + +TMP = tempfile.mkdtemp(prefix='upscale-qt-test-') +os.environ['QT_QPA_PLATFORM'] = 'offscreen' +os.environ['XDG_CONFIG_HOME'] = os.path.join(TMP, 'config') + +HERE = os.path.dirname(os.path.abspath(__file__)) +_loader = SourceFileLoader('upscale_qt', os.path.join(HERE, 'upscale-qt')) +uq = module_from_spec(spec_from_loader('upscale_qt', _loader)) +_loader.exec_module(uq) + +from PyQt6.QtGui import QColor, QImage, QImageReader # noqa: E402 + +NAMES = ['4x_detail.pth', 'RealESRGAN_x4plus.pth', + 'RealESRGAN_x4plus_anime_6B.pth'] + + +def test_match_model(): + m = uq.match_model + # the longer stem must win over its own prefix + assert m('RealESRGAN_x4plus_anime_6B.pth', NAMES) == 'RealESRGAN_x4plus_anime_6B.pth' + assert m(' realesrgan_x4plus\n', NAMES) == 'RealESRGAN_x4plus.pth' + assert m('I suggest 4x_detail.', NAMES) == '4x_detail.pth' + assert m('no idea', NAMES) is None + assert m('', []) is None + + +def test_list_models(): + d = os.path.join(TMP, 'models') + os.makedirs(d) + for f in ('b.safetensors', 'a.pth', 'notes.txt', 'c.GGUF'): + open(os.path.join(d, f), 'w').close() + assert uq.list_models(d) == ['a.pth', 'b.safetensors', 'c.GGUF'] + assert uq.list_models(os.path.join(TMP, 'missing')) == [] + + +def test_args(): + a = uq.sd_args('/m/x.pth', '/in.png', '/out.png') + assert a == ['-M', 'upscale', '--upscale-model', '/m/x.pth', + '-i', '/in.png', '-o', '/out.png'] + a = uq.llama_args('/m.gguf', '/mm.gguf', '/in.png', NAMES) + assert a[a.index('--image') + 1] == '/in.png' + prompt = a[a.index('-p') + 1] + assert all(n in prompt for n in NAMES) + assert 'noise' in prompt + + +def test_save_over(): + path = os.path.join(TMP, 'photo.JPG') # uppercase suffix on purpose + small = QImage(4, 4, QImage.Format.Format_RGB32) + small.fill(QColor('red')) + assert small.save(path, 'jpeg') + big = QImage(16, 16, QImage.Format.Format_RGB32) + big.fill(QColor('blue')) + assert uq.save_over(big, path) + assert QImage(path).width() == 16 + assert bytes(QImageReader.imageFormat(path)) == b'jpeg' + assert not os.path.exists(path + '.part') + + +if __name__ == '__main__': + for name, fn in list(globals().items()): + if name.startswith('test_'): + fn() + print('ok', name) |
