Skip to content

Commit b95ce22

Browse files
yakov-gyakov-g
committed
add tests
1 parent 09190e7 commit b95ce22

File tree

3 files changed

+158
-131
lines changed

3 files changed

+158
-131
lines changed

‎custom_model_runner/datarobot_drum/drum/utils/setup.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ def setup_options(args=None):
3232
Parsed command line options as an argparse.Namespace object.
3333
"""
3434
arg_parser = CMRunnerArgsRegistry.get_arg_parser()
35-
3635
try:
3736
import argcomplete
3837
except ImportError:

‎tests/functional/test_runtime_parameters.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -222,8 +222,8 @@ def _test_custom_model_with_runtime_params(
222222
if docker:
223223
cmd += " --docker {} --verbose ".format(docker)
224224

225-
_, stdout, _ = _exec_shell_cmd(cmd, err_msg=None, assert_if_fail=False, env=env)
226-
return stdout
225+
_, stdout, stderr = _exec_shell_cmd(cmd, err_msg=None, assert_if_fail=False, env=env)
226+
return stdout, stderr
227227

228228
@classmethod
229229
def _setup_runtime_parameters(
@@ -251,20 +251,20 @@ def _setup_runtime_parameters(
251251
def test_runtime_parameters_invalid_yaml(
252252
self, resources, tmp_path, runtime_param_values_stream, use_runtime_params_env_var
253253
):
254-
stdout = self._test_custom_model_with_runtime_params(
254+
stdout, stderr = self._test_custom_model_with_runtime_params(
255255
resources,
256256
tmp_path,
257257
runtime_param_values_stream,
258258
is_invalid_yaml=True,
259259
use_runtime_params_env_var=use_runtime_params_env_var,
260260
)
261-
assert "Invalid runtime parameter values YAML content!" in stdout
261+
assert "Invalid runtime parameter values YAML content!" in stdout + stderr
262262

263263
@pytest.mark.parametrize("use_runtime_params_env_var", [True, False])
264264
def test_runtime_parameters_missing_attr(
265265
self, resources, tmp_path, runtime_param_values_stream, use_runtime_params_env_var
266266
):
267-
stdout = self._test_custom_model_with_runtime_params(
267+
stdout, stderr = self._test_custom_model_with_runtime_params(
268268
resources,
269269
tmp_path,
270270
runtime_param_values_stream,
@@ -273,27 +273,27 @@ def test_runtime_parameters_missing_attr(
273273
)
274274
assert re.search(
275275
r".*Failed to load runtime parameter.*{\\'credentialType\\': DataError\(\\'is required\\'\)}.*",
276-
stdout,
276+
stdout + stderr,
277277
)
278278

279279
def test_runtime_parameters_boolean_invalid(
280280
self, resources, tmp_path, runtime_param_values_stream
281281
):
282-
stderr = self._test_custom_model_with_runtime_params(
282+
stdout, stderr = self._test_custom_model_with_runtime_params(
283283
resources, tmp_path, runtime_param_values_stream, bool_var_value="text"
284284
)
285285
assert re.search(
286286
r".*Failed to load runtime parameter.*value should be True or False.*",
287-
stderr,
287+
stdout + stderr,
288288
)
289289

290290
def test_runtime_parameters_numeric_invalid(
291291
self, resources, tmp_path, runtime_param_values_stream
292292
):
293-
stderr = self._test_custom_model_with_runtime_params(
293+
stdout, stderr = self._test_custom_model_with_runtime_params(
294294
resources, tmp_path, runtime_param_values_stream, numeric_var_value="text"
295295
)
296296
assert re.search(
297297
r".*Failed to load runtime parameter.*value can.*t be converted to float.*",
298-
stderr,
298+
stdout + stderr,
299299
)

‎tests/unit/datarobot_drum/drum/utils/test_setup.py‎

Lines changed: 148 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -11,123 +11,151 @@
1111
# Import the function under test
1212
from datarobot_drum.drum.utils.setup import setup_options
1313

14-
15-
def test_setup_options_default(monkeypatch):
16-
# Mock CMRunnerArgsRegistry and dependencies
17-
mock_parser = mock.Mock()
18-
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
19-
mock_parser.parse_args.return_value = mock_options
20-
21-
with mock.patch(
22-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
23-
return_value=mock_parser,
24-
), mock.patch(
25-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
26-
), mock.patch(
27-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
28-
):
29-
opts = setup_options([])
30-
assert hasattr(opts, "max_workers")
31-
assert opts.max_workers == 1 # Default to 1
32-
33-
34-
def test_setup_options_with_max_workers(monkeypatch):
35-
mock_parser = mock.Mock()
36-
mock_options = Namespace(max_workers="3", runtime_params_file=None, lazy_loading_file=None)
37-
mock_parser.parse_args.return_value = mock_options
38-
39-
with mock.patch(
40-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
41-
return_value=mock_parser,
42-
), mock.patch(
43-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
44-
), mock.patch(
45-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
46-
):
47-
opts = setup_options([])
48-
assert opts.max_workers == 3 or opts.max_workers == "3" or int(opts.max_workers) == 3
49-
50-
51-
def test_setup_options_runtime_parameters(monkeypatch):
52-
mock_parser = mock.Mock()
53-
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
54-
mock_parser.parse_args.return_value = mock_options
55-
56-
with mock.patch(
57-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
58-
return_value=mock_parser,
59-
), mock.patch(
60-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
61-
), mock.patch(
62-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
63-
), mock.patch(
64-
"datarobot_drum.RuntimeParameters.has", return_value=True
65-
), mock.patch(
66-
"datarobot_drum.RuntimeParameters.get", return_value=5
67-
):
68-
opts = setup_options([])
69-
assert opts.max_workers == 5
70-
71-
72-
def test_setup_options_runtime_params_file(monkeypatch):
73-
mock_parser = mock.Mock()
74-
mock_options = Namespace(
75-
max_workers=None, runtime_params_file="params.yaml", code_dir=".", lazy_loading_file=None
76-
)
77-
mock_parser.parse_args.return_value = mock_options
78-
79-
with mock.patch(
80-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
81-
return_value=mock_parser,
82-
), mock.patch(
83-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
84-
), mock.patch(
85-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
86-
), mock.patch(
87-
"datarobot_drum.drum.utils.setup.RuntimeParametersLoader"
88-
) as mock_loader:
89-
instance = mock_loader.return_value
90-
instance.setup_environment_variables = mock.Mock()
91-
opts = setup_options([])
92-
instance.setup_environment_variables.assert_called_once()
93-
94-
95-
def test_setup_options_lazy_loading_file(monkeypatch):
96-
mock_parser = mock.Mock()
97-
mock_options = Namespace(
98-
max_workers=None, runtime_params_file=None, lazy_loading_file="lazy.yaml"
99-
)
100-
mock_parser.parse_args.return_value = mock_options
101-
102-
with mock.patch(
103-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
104-
return_value=mock_parser,
105-
), mock.patch(
106-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
107-
), mock.patch(
108-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
109-
), mock.patch(
110-
"datarobot_drum.drum.lazy_loading.lazy_loading_handler.LazyLoadingHandler.setup_environment_variables_from_values_file"
111-
) as mock_lazy:
112-
opts = setup_options([])
113-
mock_lazy.assert_called_once_with("lazy.yaml")
114-
115-
116-
def test_setup_options_argcomplete_missing(monkeypatch, capsys):
117-
mock_parser = mock.Mock()
118-
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
119-
mock_parser.parse_args.return_value = mock_options
120-
121-
with mock.patch(
122-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
123-
return_value=mock_parser,
124-
), mock.patch(
125-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
126-
), mock.patch(
127-
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
128-
), mock.patch.dict(
129-
"sys.modules", {"argcomplete": None}
130-
):
131-
opts = setup_options([])
132-
captured = capsys.readouterr()
133-
assert "autocompletion of arguments is not supported" in captured.err
14+
import os
15+
16+
17+
class TestSetupOptions:
18+
def test_setup_options_default(self, monkeypatch):
19+
# Mock CMRunnerArgsRegistry and dependencies
20+
mock_parser = mock.Mock()
21+
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
22+
mock_parser.parse_args.return_value = mock_options
23+
24+
with mock.patch(
25+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
26+
return_value=mock_parser,
27+
), mock.patch(
28+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
29+
), mock.patch(
30+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
31+
):
32+
opts = setup_options([])
33+
assert hasattr(opts, "max_workers")
34+
assert opts.max_workers == 1 # Default to 1
35+
36+
def test_setup_options_with_max_workers(self, monkeypatch):
37+
mock_parser = mock.Mock()
38+
mock_options = Namespace(max_workers="3", runtime_params_file=None, lazy_loading_file=None)
39+
mock_parser.parse_args.return_value = mock_options
40+
41+
with mock.patch(
42+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
43+
return_value=mock_parser,
44+
), mock.patch(
45+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
46+
), mock.patch(
47+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
48+
):
49+
opts = setup_options([])
50+
assert opts.max_workers == 3 or opts.max_workers == "3" or int(opts.max_workers) == 3
51+
52+
def test_setup_options_runtime_parameters(self, monkeypatch):
53+
mock_parser = mock.Mock()
54+
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
55+
mock_parser.parse_args.return_value = mock_options
56+
57+
with mock.patch(
58+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
59+
return_value=mock_parser,
60+
), mock.patch(
61+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
62+
), mock.patch(
63+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
64+
), mock.patch(
65+
"datarobot_drum.RuntimeParameters.has", return_value=True
66+
), mock.patch(
67+
"datarobot_drum.RuntimeParameters.get", return_value=5
68+
):
69+
opts = setup_options([])
70+
assert opts.max_workers == 5
71+
72+
def test_setup_options_runtime_params_file(self, monkeypatch):
73+
mock_parser = mock.Mock()
74+
mock_options = Namespace(
75+
max_workers=None,
76+
runtime_params_file="params.yaml",
77+
code_dir=".",
78+
lazy_loading_file=None,
79+
)
80+
mock_parser.parse_args.return_value = mock_options
81+
82+
with mock.patch(
83+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
84+
return_value=mock_parser,
85+
), mock.patch(
86+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
87+
), mock.patch(
88+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
89+
), mock.patch(
90+
"datarobot_drum.drum.utils.setup.RuntimeParametersLoader"
91+
) as mock_loader:
92+
instance = mock_loader.return_value
93+
instance.setup_environment_variables = mock.Mock()
94+
opts = setup_options([])
95+
instance.setup_environment_variables.assert_called_once()
96+
97+
def test_setup_options_lazy_loading_file(self, monkeypatch):
98+
mock_parser = mock.Mock()
99+
mock_options = Namespace(
100+
max_workers=None, runtime_params_file=None, lazy_loading_file="lazy.yaml"
101+
)
102+
mock_parser.parse_args.return_value = mock_options
103+
104+
with mock.patch(
105+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
106+
return_value=mock_parser,
107+
), mock.patch(
108+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
109+
), mock.patch(
110+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
111+
), mock.patch(
112+
"datarobot_drum.drum.lazy_loading.lazy_loading_handler.LazyLoadingHandler.setup_environment_variables_from_values_file"
113+
) as mock_lazy:
114+
opts = setup_options([])
115+
mock_lazy.assert_called_once_with("lazy.yaml")
116+
117+
def test_setup_options_argcomplete_missing(self, monkeypatch, capsys):
118+
mock_parser = mock.Mock()
119+
mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None)
120+
mock_parser.parse_args.return_value = mock_options
121+
122+
with mock.patch(
123+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
124+
return_value=mock_parser,
125+
), mock.patch(
126+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
127+
), mock.patch(
128+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
129+
), mock.patch.dict(
130+
"sys.modules", {"argcomplete": None}
131+
):
132+
opts = setup_options([])
133+
captured = capsys.readouterr()
134+
assert "autocompletion of arguments is not supported" in captured.err
135+
136+
def test_setup_options_address_env_var(self, monkeypatch):
137+
mock_parser = mock.Mock()
138+
test_address = "127.0.0.1:9999"
139+
mock_options = Namespace(
140+
max_workers=None, runtime_params_file=None, lazy_loading_file=None, address=test_address
141+
)
142+
mock_parser.parse_args.return_value = mock_options
143+
144+
# Remove ADDRESS if it exists to avoid side effects
145+
if "ADDRESS" in os.environ:
146+
del os.environ["ADDRESS"]
147+
148+
with mock.patch(
149+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser",
150+
return_value=mock_parser,
151+
), mock.patch(
152+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars"
153+
), mock.patch(
154+
"datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options"
155+
):
156+
opts = setup_options([])
157+
assert os.environ["ADDRESS"] == test_address
158+
159+
# Clean up
160+
if "ADDRESS" in os.environ:
161+
del os.environ["ADDRESS"]

0 commit comments

Comments
 (0)