1
0
Fork 0
colorclass/tests/test_windows.py

430 lines
16 KiB
Python
Raw Permalink Normal View History

"""Test Windows methods."""
from __future__ import print_function
import ctypes
import sys
from textwrap import dedent
try:
from StringIO import StringIO
except ImportError:
from io import StringIO
import pytest
from colorclass import windows
from colorclass.codes import ANSICodeMapping
from colorclass.color import Color
from tests.conftest import PROJECT_ROOT
from tests.screenshot import RunNewConsole, screenshot_until_match
class MockKernel32(object):
"""Mock kernel32."""
def __init__(self, stderr=windows.INVALID_HANDLE_VALUE, stdout=windows.INVALID_HANDLE_VALUE, set_mode=0x0):
"""Constructor."""
self.set_mode = set_mode
self.stderr = stderr
self.stdout = stdout
self.wAttributes = 7
def GetConsoleMode(self, _, dword_pointer): # noqa
"""Mock GetConsoleMode.
:param _: Unused handle.
:param dword_pointer: ctypes.byref(lpdword) return value.
"""
ulong_ptr = ctypes.POINTER(ctypes.c_ulong)
dword = ctypes.cast(dword_pointer, ulong_ptr).contents # Dereference pointer.
dword.value = self.set_mode
return 1
def GetConsoleScreenBufferInfo(self, _, csbi_pointer): # noqa
"""Mock GetConsoleScreenBufferInfo.
:param _: Unused handle.
:param csbi_pointer: ctypes.byref(csbi) return value.
"""
struct_ptr = ctypes.POINTER(windows.ConsoleScreenBufferInfo)
csbi = ctypes.cast(csbi_pointer, struct_ptr).contents # Dereference pointer.
csbi.wAttributes = self.wAttributes
return 1
def GetStdHandle(self, handle): # noqa
"""Mock GetStdHandle.
:param int handle: STD_ERROR_HANDLE or STD_OUTPUT_HANDLE.
"""
return self.stderr if handle == windows.STD_ERROR_HANDLE else self.stdout
def SetConsoleTextAttribute(self, _, color_code): # noqa
"""Mock SetConsoleTextAttribute.
:param _: Unused handle.
:param int color_code: Merged color code to set.
"""
self.wAttributes = color_code
return 1
class MockSys(object):
"""Mock sys standard library module."""
def __init__(self, stderr=None, stdout=None):
"""Constructor."""
self.stderr = stderr or type('', (), {})
self.stdout = stdout or type('', (), {})
@pytest.mark.skipif(str(not windows.IS_WINDOWS))
def test_init_kernel32_unique():
"""Make sure function doesn't override other LibraryLoaders."""
k32_a = ctypes.LibraryLoader(ctypes.WinDLL).kernel32
k32_a.GetStdHandle.argtypes = [ctypes.c_void_p]
k32_a.GetStdHandle.restype = ctypes.c_ulong
k32_b, stderr_b, stdout_b = windows.init_kernel32()
k32_c = ctypes.LibraryLoader(ctypes.WinDLL).kernel32
k32_c.GetStdHandle.argtypes = [ctypes.c_long]
k32_c.GetStdHandle.restype = ctypes.c_short
k32_d, stderr_d, stdout_d = windows.init_kernel32()
# Verify external.
assert k32_a.GetStdHandle.argtypes == [ctypes.c_void_p]
assert k32_a.GetStdHandle.restype == ctypes.c_ulong
assert k32_c.GetStdHandle.argtypes == [ctypes.c_long]
assert k32_c.GetStdHandle.restype == ctypes.c_short
# Verify ours.
assert k32_b.GetStdHandle.argtypes == [ctypes.c_ulong]
assert k32_b.GetStdHandle.restype == ctypes.c_void_p
assert k32_d.GetStdHandle.argtypes == [ctypes.c_ulong]
assert k32_d.GetStdHandle.restype == ctypes.c_void_p
assert stderr_b == stderr_d
assert stdout_b == stdout_d
@pytest.mark.parametrize('stderr_invalid', [False, True])
@pytest.mark.parametrize('stdout_invalid', [False, True])
def test_init_kernel32_valid_handle(monkeypatch, stderr_invalid, stdout_invalid):
"""Test valid/invalid handle handling.
:param monkeypatch: pytest fixture.
:param bool stderr_invalid: Mock stderr is valid.
:param bool stdout_invalid: Mock stdout is valid.
"""
mock_sys = MockSys()
monkeypatch.setattr(windows, 'sys', mock_sys)
if stderr_invalid:
setattr(mock_sys.stderr, '_original_stream', True)
if stdout_invalid:
setattr(mock_sys.stdout, '_original_stream', True)
stderr, stdout = windows.init_kernel32(MockKernel32(stderr=100, stdout=200))[1:]
if stderr_invalid and stdout_invalid:
assert stderr == windows.INVALID_HANDLE_VALUE
assert stdout == windows.INVALID_HANDLE_VALUE
elif stdout_invalid:
assert stderr == 100
assert stdout == windows.INVALID_HANDLE_VALUE
elif stderr_invalid:
assert stderr == windows.INVALID_HANDLE_VALUE
assert stdout == 200
else:
assert stderr == 100
assert stdout == 200
def test_get_console_info():
"""Test function."""
# Test error.
if windows.IS_WINDOWS:
with pytest.raises(OSError):
windows.get_console_info(windows.init_kernel32()[0], windows.INVALID_HANDLE_VALUE)
# Test no error with mock methods.
kernel32 = MockKernel32()
fg_color, bg_color, native_ansi = windows.get_console_info(kernel32, windows.INVALID_HANDLE_VALUE)
assert fg_color == 7
assert bg_color == 0
assert native_ansi is False
# Test different console modes.
for not_native in (0x0, 0x1, 0x2, 0x1 | 0x2):
kernel32.set_mode = not_native
assert not windows.get_console_info(kernel32, windows.INVALID_HANDLE_VALUE)[-1]
for native in (i | 0x4 for i in (0x0, 0x1, 0x2, 0x1 | 0x2)):
kernel32.set_mode = native
assert windows.get_console_info(kernel32, windows.INVALID_HANDLE_VALUE)[-1]
@pytest.mark.parametrize('stderr', [1, windows.INVALID_HANDLE_VALUE])
@pytest.mark.parametrize('stdout', [2, windows.INVALID_HANDLE_VALUE])
def test_bg_color_native_ansi(stderr, stdout):
"""Test function.
:param int stderr: Value of parameter.
:param int stdout: Value of parameter.
"""
kernel32 = MockKernel32(set_mode=0x4)
kernel32.wAttributes = 240
actual = windows.bg_color_native_ansi(kernel32, stderr, stdout)
if stderr == windows.INVALID_HANDLE_VALUE and stdout == windows.INVALID_HANDLE_VALUE:
expected = 0, False
else:
expected = 240, True
assert actual == expected
def test_windows_stream():
"""Test class."""
# Test error.
if windows.IS_WINDOWS:
stream = windows.WindowsStream(windows.init_kernel32()[0], windows.INVALID_HANDLE_VALUE, StringIO())
assert stream.colors == (windows.WINDOWS_CODES['white'], windows.WINDOWS_CODES['black'])
stream.colors = windows.WINDOWS_CODES['red'] | windows.WINDOWS_CODES['bgblue'] # No exception, just ignore.
assert stream.colors == (windows.WINDOWS_CODES['white'], windows.WINDOWS_CODES['black'])
# Test __getattr__() and color resetting.
original_stream = StringIO()
stream = windows.WindowsStream(MockKernel32(), windows.INVALID_HANDLE_VALUE, original_stream)
assert stream.writelines == original_stream.writelines # Test __getattr__().
assert stream.colors == (windows.WINDOWS_CODES['white'], windows.WINDOWS_CODES['black'])
stream.colors = windows.WINDOWS_CODES['red'] | windows.WINDOWS_CODES['bgblue']
assert stream.colors == (windows.WINDOWS_CODES['red'], windows.WINDOWS_CODES['bgblue'])
stream.colors = None # Resets colors to original.
assert stream.colors == (windows.WINDOWS_CODES['white'], windows.WINDOWS_CODES['black'])
# Test special negative codes.
stream.colors = windows.WINDOWS_CODES['red'] | windows.WINDOWS_CODES['bgblue']
stream.colors = windows.WINDOWS_CODES['/fg']
assert stream.colors == (windows.WINDOWS_CODES['white'], windows.WINDOWS_CODES['bgblue'])
stream.colors = windows.WINDOWS_CODES['red'] | windows.WINDOWS_CODES['bgblue']
stream.colors = windows.WINDOWS_CODES['/bg']
assert stream.colors == (windows.WINDOWS_CODES['red'], windows.WINDOWS_CODES['black'])
stream.colors = windows.WINDOWS_CODES['red'] | windows.WINDOWS_CODES['bgblue']
stream.colors = windows.WINDOWS_CODES['bgblack']
assert stream.colors == (windows.WINDOWS_CODES['red'], windows.WINDOWS_CODES['black'])
# Test write.
stream.write(Color('{/all}A{red}B{bgblue}C'))
original_stream.seek(0)
assert original_stream.read() == 'ABC'
assert stream.colors == (windows.WINDOWS_CODES['red'], windows.WINDOWS_CODES['bgblue'])
# Test ignore invalid code.
original_stream.seek(0)
original_stream.truncate()
stream.write('\x1b[0mA\x1b[31mB\x1b[44;999mC')
original_stream.seek(0)
assert original_stream.read() == 'ABC'
assert stream.colors == (windows.WINDOWS_CODES['red'], windows.WINDOWS_CODES['bgblue'])
@pytest.mark.skipif(str(windows.IS_WINDOWS))
def test_windows_nix():
"""Test enable/disable on non-Windows platforms."""
with windows.Windows():
assert not windows.Windows.is_enabled()
assert not hasattr(sys.stderr, '_original_stream')
assert not hasattr(sys.stdout, '_original_stream')
assert not windows.Windows.is_enabled()
assert not hasattr(sys.stderr, '_original_stream')
assert not hasattr(sys.stdout, '_original_stream')
def test_windows_auto_colors(monkeypatch):
"""Test Windows class with/out valid_handle and with/out auto_colors. Don't replace streams.
:param monkeypatch: pytest fixture.
"""
mock_sys = MockSys()
monkeypatch.setattr(windows, 'atexit', type('', (), {'register': staticmethod(lambda _: 0 / 0)}))
monkeypatch.setattr(windows, 'IS_WINDOWS', True)
monkeypatch.setattr(windows, 'sys', mock_sys)
monkeypatch.setattr(ANSICodeMapping, 'LIGHT_BACKGROUND', None)
# Test no valid handles.
kernel32 = MockKernel32()
monkeypatch.setattr(windows, 'init_kernel32', lambda: (kernel32, -1, -1))
assert not windows.Windows.enable()
assert not windows.Windows.is_enabled()
assert not hasattr(mock_sys.stderr, '_original_stream')
assert not hasattr(mock_sys.stdout, '_original_stream')
assert ANSICodeMapping.LIGHT_BACKGROUND is None
# Test auto colors dark background.
kernel32.set_mode = 0x4 # Enable native ANSI to have Windows skip replacing streams.
monkeypatch.setattr(windows, 'init_kernel32', lambda: (kernel32, 1, 2))
assert not windows.Windows.enable(auto_colors=True)
assert not windows.Windows.is_enabled()
assert not hasattr(mock_sys.stderr, '_original_stream')
assert not hasattr(mock_sys.stdout, '_original_stream')
assert ANSICodeMapping.LIGHT_BACKGROUND is False
# Test auto colors light background.
kernel32.wAttributes = 240
assert not windows.Windows.enable(auto_colors=True)
assert not windows.Windows.is_enabled()
assert not hasattr(mock_sys.stderr, '_original_stream')
assert not hasattr(mock_sys.stdout, '_original_stream')
assert ANSICodeMapping.LIGHT_BACKGROUND is True
@pytest.mark.parametrize('valid', ['stderr', 'stdout', 'both'])
def test_windows_replace_streams(monkeypatch, tmpdir, valid):
"""Test Windows class stdout and stderr replacement.
:param monkeypatch: pytest fixture.
:param tmpdir: pytest fixture.
:param str valid: Which mock stream(s) should be valid.
"""
ac = list() # atexit called.
mock_sys = MockSys(stderr=tmpdir.join('stderr').open(mode='wb'), stdout=tmpdir.join('stdout').open(mode='wb'))
monkeypatch.setattr(windows, 'atexit', type('', (), {'register': staticmethod(lambda _: ac.append(1))}))
monkeypatch.setattr(windows, 'IS_WINDOWS', True)
monkeypatch.setattr(windows, 'sys', mock_sys)
# Mock init_kernel32.
stderr = 1 if valid in ('stderr', 'both') else windows.INVALID_HANDLE_VALUE
stdout = 2 if valid in ('stdout', 'both') else windows.INVALID_HANDLE_VALUE
monkeypatch.setattr(windows, 'init_kernel32', lambda: (MockKernel32(), stderr, stdout))
# Test.
assert windows.Windows.enable(reset_atexit=True)
assert windows.Windows.is_enabled()
assert len(ac) == 1
if stderr != windows.INVALID_HANDLE_VALUE:
assert hasattr(mock_sys.stderr, '_original_stream')
else:
assert not hasattr(mock_sys.stderr, '_original_stream')
if stdout != windows.INVALID_HANDLE_VALUE:
assert hasattr(mock_sys.stdout, '_original_stream')
else:
assert not hasattr(mock_sys.stdout, '_original_stream')
# Test multiple disable.
assert windows.Windows.disable()
assert not windows.Windows.is_enabled()
assert not windows.Windows.disable()
assert not windows.Windows.is_enabled()
# Test context manager.
with windows.Windows():
assert windows.Windows.is_enabled()
assert not windows.Windows.is_enabled()
@pytest.mark.skipif(str(not windows.IS_WINDOWS))
def test_enable_disable(tmpdir):
"""Test enabling, disabling, repeat. Make sure colors still work.
:param tmpdir: pytest fixture.
"""
screenshot = PROJECT_ROOT.join('test_windows_test_enable_disable.png')
if screenshot.check():
screenshot.remove()
script = tmpdir.join('script.py')
command = [sys.executable, str(script)]
script.write(dedent("""\
from __future__ import print_function
import os, time
from colorclass import Color, Windows
with Windows(auto_colors=True):
print(Color('{autored}Red{/autored}'))
print('Red')
with Windows(auto_colors=True):
print(Color('{autored}Red{/autored}'))
print('Red')
stop_after = time.time() + 20
while not os.path.exists(r'%s') and time.time() < stop_after:
time.sleep(0.5)
""") % str(screenshot))
# Setup expected.
with_colors = [str(p) for p in PROJECT_ROOT.join('tests').listdir('sub_red_light_fg_*.bmp')]
sans_colors = [str(p) for p in PROJECT_ROOT.join('tests').listdir('sub_red_sans_*.bmp')]
assert with_colors
assert sans_colors
# Run.
with RunNewConsole(command) as gen:
screenshot_until_match(str(screenshot), 15, with_colors, 2, gen)
screenshot_until_match(str(screenshot), 15, sans_colors, 2, gen)
@pytest.mark.skipif(str(not windows.IS_WINDOWS))
def test_box_characters(tmpdir):
"""Test for unicode errors with special characters.
:param tmpdir: pytest fixture.
"""
screenshot = PROJECT_ROOT.join('test_windows_test_box_characters.png')
if screenshot.check():
screenshot.remove()
script = tmpdir.join('script.py')
command = [sys.executable, str(script)]
script.write(dedent("""\
from __future__ import print_function
import os, time
from colorclass import Color, Windows
Windows.enable(auto_colors=True)
chars = [
'+', '-', '|',
b'\\xb3'.decode('ibm437'),
b'\\xb4'.decode('ibm437'),
b'\\xb9'.decode('ibm437'),
b'\\xba'.decode('ibm437'),
b'\\xbb'.decode('ibm437'),
b'\\xbc'.decode('ibm437'),
b'\\xbf'.decode('ibm437'),
b'\\xc0'.decode('ibm437'),
b'\\xc1'.decode('ibm437'),
b'\\xc2'.decode('ibm437'),
b'\\xc3'.decode('ibm437'),
b'\\xc4'.decode('ibm437'),
b'\\xc5'.decode('ibm437'),
b'\\xc8'.decode('ibm437'),
b'\\xc9'.decode('ibm437'),
b'\\xca'.decode('ibm437'),
b'\\xcb'.decode('ibm437'),
b'\\xcc'.decode('ibm437'),
b'\\xcd'.decode('ibm437'),
b'\\xce'.decode('ibm437'),
b'\\xd9'.decode('ibm437'),
b'\\xda'.decode('ibm437'),
]
for c in chars:
print(c, end='')
print()
for c in chars:
print(Color.green(c, auto=True), end='')
print()
stop_after = time.time() + 20
while not os.path.exists(r'%s') and time.time() < stop_after:
time.sleep(0.5)
""") % str(screenshot))
# Setup expected.
with_colors = [str(p) for p in PROJECT_ROOT.join('tests').listdir('sub_box_green_*.bmp')]
sans_colors = [str(p) for p in PROJECT_ROOT.join('tests').listdir('sub_box_sans_*.bmp')]
assert with_colors
assert sans_colors
# Run.
with RunNewConsole(command) as gen:
screenshot_until_match(str(screenshot), 15, with_colors, 1, gen)
screenshot_until_match(str(screenshot), 15, sans_colors, 1, gen)