mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 11:08:44 +00:00
170 lines
8.2 KiB
Python
170 lines
8.2 KiB
Python
import hashlib
|
|
import io
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import ExitStack
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from utils import update_check, windows_update
|
|
|
|
|
|
class AssetNameTest(unittest.TestCase):
|
|
def test_keeps_the_running_build_even_when_renamed_or_emulated(self):
|
|
builds = (
|
|
('win-amd64', (3, 11), 'TgWsProxy_windows.exe'),
|
|
('win-amd64', (3, 8), 'TgWsProxy_windows_7_64bit.exe'),
|
|
('win32', (3, 8), 'TgWsProxy_windows_7_32bit.exe'),
|
|
('win-arm64', (3, 11), 'TgWsProxy_windows_arm64.exe'),
|
|
('unknown', (3, 11), None),
|
|
)
|
|
for platform, version, expected in builds:
|
|
with self.subTest(platform=platform, version=version), \
|
|
mock.patch.object(windows_update.sysconfig, 'get_platform', return_value=platform), \
|
|
mock.patch.object(sys, 'version_info', version), \
|
|
mock.patch.object(sys, 'executable', 'renamed.exe'):
|
|
self.assertEqual(windows_update.asset_name(), expected)
|
|
|
|
|
|
class ApplyUpdateTest(unittest.TestCase):
|
|
def setUp(self):
|
|
stack = ExitStack()
|
|
self.addCleanup(stack.close)
|
|
directory = stack.enter_context(tempfile.TemporaryDirectory())
|
|
self.exe = Path(directory) / 'renamed app.exe'
|
|
self.exe.write_bytes(b'old')
|
|
self.backup = self.exe.with_name('renamed app_oldtgws.exe')
|
|
self.opener = mock.Mock()
|
|
self.opener.open.return_value = io.BytesIO(b'new')
|
|
self.asset = {'url': 'https://example.invalid/new.exe',
|
|
'digest': 'sha256:' + hashlib.sha256(b'new').hexdigest()}
|
|
stack.enter_context(mock.patch.object(sys, 'executable', str(self.exe)))
|
|
stack.enter_context(mock.patch.object(sys, 'argv', [str(self.exe), '--portable']))
|
|
stack.enter_context(mock.patch.object(sys, '_MEIPASS', None, create=True))
|
|
stack.enter_context(mock.patch.object(update_check, 'build_github_opener', return_value=self.opener))
|
|
stack.enter_context(mock.patch.object(windows_update.subprocess, 'DETACHED_PROCESS', 8, create=True))
|
|
stack.enter_context(mock.patch.object(windows_update.subprocess, 'CREATE_NEW_PROCESS_GROUP', 512, create=True))
|
|
self.launch = stack.enter_context(mock.patch.object(windows_update.subprocess, 'Popen'))
|
|
self.prepare = mock.Mock()
|
|
|
|
def test_verified_update_keeps_legacy_backup_and_launch_arguments(self):
|
|
def prepare():
|
|
self.assertEqual(self.exe.read_bytes(), b'new')
|
|
self.assertEqual(self.backup.read_bytes(), b'old')
|
|
self.launch.assert_not_called()
|
|
|
|
windows_update.apply_update(self.asset, prepare)
|
|
self.assertEqual(self.exe.read_bytes(), b'new')
|
|
self.assertEqual(self.backup.read_bytes(), b'old')
|
|
self.assertEqual(set(self.exe.parent.iterdir()), {self.exe, self.backup})
|
|
self.assertEqual(self.launch.call_args[0][0], [str(self.exe), '--portable'])
|
|
self.assertEqual(self.launch.call_args[1]['creationflags'], 520)
|
|
|
|
def test_missing_or_corrupt_checksum_does_not_replace_or_launch(self):
|
|
for digest in ('', 'sha256:' + '0' * 64):
|
|
with self.subTest(digest=digest):
|
|
self.asset['digest'] = digest
|
|
with self.assertRaises((OSError, ValueError)):
|
|
windows_update.apply_update(self.asset, self.prepare)
|
|
self.assertEqual(self.exe.read_bytes(), b'old')
|
|
self.assertEqual(list(self.exe.parent.iterdir()), [self.exe])
|
|
self.prepare.assert_not_called()
|
|
self.launch.assert_not_called()
|
|
|
|
def test_launch_failure_restores_the_old_executable(self):
|
|
self.launch.side_effect = OSError('launch failed')
|
|
with self.assertRaisesRegex(OSError, 'launch failed'):
|
|
windows_update.apply_update(self.asset, self.prepare)
|
|
self.prepare.assert_called_once_with()
|
|
self.assertEqual(self.exe.read_bytes(), b'old')
|
|
self.assertEqual(list(self.exe.parent.iterdir()), [self.exe])
|
|
|
|
def test_rename_failure_leaves_current_executable_and_cleans_download(self):
|
|
with mock.patch.object(Path, 'rename', side_effect=PermissionError('in use')):
|
|
with self.assertRaises(PermissionError):
|
|
windows_update.apply_update(self.asset, self.prepare)
|
|
self.assertEqual(self.exe.read_bytes(), b'old')
|
|
self.assertEqual(list(self.exe.parent.iterdir()), [self.exe])
|
|
self.prepare.assert_not_called()
|
|
self.launch.assert_not_called()
|
|
|
|
def test_replacement_failure_rolls_back_and_cleans_download(self):
|
|
replace = Path.replace
|
|
|
|
def fail_download_move(path, target):
|
|
if path.suffix == '.tmp':
|
|
raise PermissionError('move failed')
|
|
return replace(path, target)
|
|
|
|
with mock.patch.object(Path, 'replace', fail_download_move):
|
|
with self.assertRaisesRegex(PermissionError, 'move failed'):
|
|
windows_update.apply_update(self.asset, self.prepare)
|
|
self.assertEqual(self.exe.read_bytes(), b'old')
|
|
self.assertEqual(list(self.exe.parent.iterdir()), [self.exe])
|
|
self.prepare.assert_not_called()
|
|
self.launch.assert_not_called()
|
|
|
|
def test_restart_uses_a_fresh_bundle_and_resets_dll_directory(self):
|
|
dll = mock.Mock()
|
|
with mock.patch.object(sys, '_MEIPASS', 'old-bundle'), \
|
|
mock.patch.object(windows_update.ctypes, 'windll', dll, create=True):
|
|
windows_update.apply_update(self.asset, self.prepare)
|
|
self.assertEqual(self.launch.call_args[1]['env']['PYINSTALLER_RESET_ENVIRONMENT'], '1')
|
|
calls = dll.kernel32.SetDllDirectoryW.call_args_list
|
|
self.assertEqual(calls[0], mock.call(None))
|
|
self.assertEqual(calls[1].args[0].value, 'old-bundle')
|
|
|
|
|
|
@unittest.skipUnless(sys.platform == 'win32', 'Windows tray integration')
|
|
class UpdateFlowTest(unittest.TestCase):
|
|
def setUp(self):
|
|
import windows
|
|
|
|
self.windows = windows
|
|
|
|
def test_fallback_dialog_actually_installs_the_selected_asset(self):
|
|
app = self.windows
|
|
asset = {'url': 'https://example.invalid/new.exe', 'digest': 'sha256:' + '0' * 64}
|
|
with mock.patch.object(app, 'ctk', None), \
|
|
mock.patch.object(app._u32, 'MessageBoxW', return_value=app._IDYES), \
|
|
mock.patch.object(app, '_perform_update') as update:
|
|
app.update_ctk_form('update', asset=asset)
|
|
update.assert_called_once_with(asset)
|
|
|
|
def test_failed_launch_recovers_proxy_and_lock_instead_of_exiting(self):
|
|
app = self.windows
|
|
|
|
def fail(asset, before_restart):
|
|
before_restart()
|
|
raise OSError('launch failed')
|
|
|
|
with ExitStack() as stack:
|
|
mocks = {name: stack.enter_context(mock.patch.object(app, name)) for name in (
|
|
'apply_update', 'stop_proxy', 'release_lock', '_release_win_mutex',
|
|
'_acquire_win_mutex', 'acquire_lock', 'start_proxy', '_show_error',
|
|
)}
|
|
exit_process = stack.enter_context(mock.patch.object(app.os, '_exit'))
|
|
mocks['apply_update'].side_effect = fail
|
|
mocks['_acquire_win_mutex'].return_value = None
|
|
app._perform_update({'url': 'asset', 'digest': 'digest'})
|
|
mocks['stop_proxy'].assert_called_once_with()
|
|
mocks['release_lock'].assert_called_once_with()
|
|
mocks['_release_win_mutex'].assert_called_once_with()
|
|
mocks['acquire_lock'].assert_called_once_with()
|
|
mocks['start_proxy'].assert_called_once_with(app._config, mocks['_show_error'])
|
|
mocks['_show_error'].assert_called_once()
|
|
exit_process.assert_not_called()
|
|
|
|
def test_download_failure_leaves_proxy_and_locks_running(self):
|
|
app = self.windows
|
|
with mock.patch.object(app, 'apply_update', side_effect=OSError('checksum mismatch')), \
|
|
mock.patch.object(app, 'stop_proxy') as stop, \
|
|
mock.patch.object(app, 'release_lock') as unlock, \
|
|
mock.patch.object(app, '_show_error'), \
|
|
mock.patch.object(app.os, '_exit') as exit_process:
|
|
app._perform_update({})
|
|
stop.assert_not_called()
|
|
unlock.assert_not_called()
|
|
exit_process.assert_not_called()
|