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
|
# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
from unittest import mock
import logging
import os
import re
import shutil
import sys
import tempfile
from io import StringIO
from knack.cli import CLI, CLICommandsLoader, CommandInvoker
from knack.log import CLI_LOGGER_NAME
TEMP_FOLDER_NAME = "knack_temp"
def redirect_io(func):
original_stderr = sys.stderr
original_stdout = sys.stdout
def wrapper(self):
# Ensure a clean startup - no log handlers
root_logger = logging.getLogger()
cli_logger = logging.getLogger(CLI_LOGGER_NAME)
root_logger.handlers.clear()
cli_logger.handlers.clear()
sys.stdout = sys.stderr = self.io = StringIO()
func(self)
self.io.close()
sys.stdout = original_stdout
sys.stderr = original_stderr
# Remove the handlers added by CLI, so that the next invoke call init them again with the new stderr
# Otherwise, the handlers will write to a closed StringIO from a preview test
root_logger.handlers.clear()
cli_logger.handlers.clear()
return wrapper
def disable_color(func):
def wrapper(self):
self.cli_ctx.enable_color = False
func(self)
self.cli_ctx.enable_color = True
return wrapper
def _remove_control_sequence(string):
return re.sub(r'\x1b[^m]+m', '', string)
def _remove_whitespace(string):
return re.sub(r'\s', '', string)
def assert_in_multi_line(sub_string, string):
# assert sub_string is in string, with all whitespaces, line breaks and control sequences ignored
assert _remove_whitespace(sub_string) in _remove_control_sequence(_remove_whitespace(string))
class MockContext(CLI):
def __init__(self):
super().__init__(config_dir=new_temp_folder())
loader = CLICommandsLoader(cli_ctx=self)
invocation = mock.MagicMock(spec=CommandInvoker)
invocation.data = {}
setattr(self, 'commands_loader', loader)
setattr(self, 'invocation', invocation)
class DummyCLI(CLI):
def get_cli_version(self):
return '0.1.0'
def __init__(self, **kwargs):
kwargs['config_dir'] = new_temp_folder()
super().__init__(**kwargs)
# Force to enable color
self.enable_color = True
def new_temp_folder():
temp_dir = os.path.join(tempfile.gettempdir(), TEMP_FOLDER_NAME)
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir)
os.mkdir(temp_dir)
return temp_dir
|
