1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
# 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
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_descriptions():
d = os.path.join(TMP, 'models')
with open(os.path.join(d, 'models.txt'), 'w') as f:
f.write('# comment = ignored\n'
'a.pth = Clean photos: people, nature\n'
'\n'
'no separator here\n'
' c.GGUF=Anime \n')
desc = uq.read_descriptions(d)
assert desc == {'a.pth': 'Clean photos: people, nature', 'c.GGUF': 'Anime'}
assert uq.model_label('a.pth', desc) == 'Clean photos: people, nature [a]'
assert uq.model_label('b.safetensors', desc) == 'b.safetensors'
assert uq.read_descriptions(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']
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')
from PyQt6.QtWidgets import QApplication # noqa: E402
APP = QApplication.instance() or QApplication([])
def test_compare_view():
before = QImage(10, 10, QImage.Format.Format_RGB32)
before.fill(QColor('red'))
after = QImage(40, 40, QImage.Format.Format_RGB32)
after.fill(QColor('blue'))
v = uq.CompareView()
v.resize(400, 400) # at or above the minimum size
v.set_images(before, after)
shot = v.grab().toImage()
# left of the divider shows before, right shows after
assert shot.pixelColor(100, 200) == QColor('red')
assert shot.pixelColor(300, 200) == QColor('blue')
def test_main_window():
w = uq.Main()
w.cfg.setValue('upscalers_dir', os.path.join(TMP, 'models'))
w.reload_models()
assert not w.a_up.isEnabled() # no image yet
path = os.path.join(TMP, 'in.png')
img = QImage(8, 6, QImage.Format.Format_RGB32)
img.fill(QColor('green'))
img.save(path)
w.load(path)
assert len(w.jobs) == 1 and w.jobs[0].image.size() == img.size()
assert os.path.exists(w.jobs[0].input)
assert w.a_up.isEnabled() and not w.a_save.isEnabled()
assert w.models.itemText(0) == 'Clean photos: people, nature [a]'
assert w.models.currentData() == 'a.pth'
assert w.a_clear.isEnabled()
temp_in = w.jobs[0].input
w.clear()
assert w.jobs == [] and w.view.before is None and not os.path.exists(temp_in)
assert not w.a_clear.isEnabled() and not w.a_up.isEnabled()
def test_quit_asks_when_unsaved():
from PyQt6.QtWidgets import QMessageBox
w = uq.Main()
assert w.a_quit.shortcut().toString() == 'Ctrl+Q'
path = os.path.join(TMP, 'q.png')
QImage(4, 4, QImage.Format.Format_RGB32).save(path)
w.load(path)
w.jobs[0].result = QImage(16, 16, QImage.Format.Format_RGB32)
answers = []
real = QMessageBox.question
try:
QMessageBox.question = lambda *a: answers.append(a) or QMessageBox.StandardButton.No
assert not w.close() and len(answers) == 1 # declined, stays open
w.jobs[0].saved = True
assert w.close() and len(answers) == 1 # saved, no question
w.what = 'upscale' # pretend sd-cli runs
assert not w.close() and len(answers) == 2 # declined, stays open
QMessageBox.question = lambda *a: answers.append(a) or QMessageBox.StandardButton.Yes
assert w.close() and len(answers) == 3
finally:
QMessageBox.question = real
if __name__ == '__main__':
for name, fn in list(globals().items()):
if name.startswith('test_'):
fn()
print('ok', name)
|
