Skip to content

Commit bd84c02

Browse files
yakov-gyakov-g
committed
add tests
1 parent 58d305d commit bd84c02

File tree

2 files changed

+148
-121
lines changed

2 files changed

+148
-121
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/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)