|
7 | 7 |
|
8 | 8 | import pytest
|
9 | 9 |
|
10 |
| -from vllm.utils import deprecate_kwargs, get_open_port, merge_async_iterators |
| 10 | +from vllm.utils import (FlexibleArgumentParser, deprecate_kwargs, |
| 11 | + get_open_port, merge_async_iterators) |
11 | 12 |
|
12 | 13 | from .utils import error_on_warning
|
13 | 14 |
|
@@ -130,3 +131,61 @@ def test_get_open_port():
|
130 | 131 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
|
131 | 132 | s3.bind(("localhost", get_open_port()))
|
132 | 133 | os.environ.pop("VLLM_PORT")
|
| 134 | + |
| 135 | + |
| 136 | +# Tests for FlexibleArgumentParser |
| 137 | +@pytest.fixture |
| 138 | +def parser(): |
| 139 | + parser = FlexibleArgumentParser() |
| 140 | + parser.add_argument('--image-input-type', |
| 141 | + choices=['pixel_values', 'image_features']) |
| 142 | + parser.add_argument('--model-name') |
| 143 | + parser.add_argument('--batch-size', type=int) |
| 144 | + parser.add_argument('--enable-feature', action='store_true') |
| 145 | + return parser |
| 146 | + |
| 147 | + |
| 148 | +def test_underscore_to_dash(parser): |
| 149 | + args = parser.parse_args(['--image_input_type', 'pixel_values']) |
| 150 | + assert args.image_input_type == 'pixel_values' |
| 151 | + |
| 152 | + |
| 153 | +def test_mixed_usage(parser): |
| 154 | + args = parser.parse_args([ |
| 155 | + '--image_input_type', 'image_features', '--model-name', |
| 156 | + 'facebook/opt-125m' |
| 157 | + ]) |
| 158 | + assert args.image_input_type == 'image_features' |
| 159 | + assert args.model_name == 'facebook/opt-125m' |
| 160 | + |
| 161 | + |
| 162 | +def test_with_equals_sign(parser): |
| 163 | + args = parser.parse_args( |
| 164 | + ['--image_input_type=pixel_values', '--model-name=facebook/opt-125m']) |
| 165 | + assert args.image_input_type == 'pixel_values' |
| 166 | + assert args.model_name == 'facebook/opt-125m' |
| 167 | + |
| 168 | + |
| 169 | +def test_with_int_value(parser): |
| 170 | + args = parser.parse_args(['--batch_size', '32']) |
| 171 | + assert args.batch_size == 32 |
| 172 | + args = parser.parse_args(['--batch-size', '32']) |
| 173 | + assert args.batch_size == 32 |
| 174 | + |
| 175 | + |
| 176 | +def test_with_bool_flag(parser): |
| 177 | + args = parser.parse_args(['--enable_feature']) |
| 178 | + assert args.enable_feature is True |
| 179 | + args = parser.parse_args(['--enable-feature']) |
| 180 | + assert args.enable_feature is True |
| 181 | + |
| 182 | + |
| 183 | +def test_invalid_choice(parser): |
| 184 | + with pytest.raises(SystemExit): |
| 185 | + parser.parse_args(['--image_input_type', 'invalid_choice']) |
| 186 | + |
| 187 | + |
| 188 | +def test_missing_required_argument(parser): |
| 189 | + parser.add_argument('--required-arg', required=True) |
| 190 | + with pytest.raises(SystemExit): |
| 191 | + parser.parse_args([]) |
0 commit comments