test_prompt_validation.py 907 Bytes
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
5
import pytest

from vllm import LLM
6
from vllm.config import LoadFormat
7
8


Joe Runde's avatar
Joe Runde committed
9
10
11
12
13
14
15
16
@pytest.fixture(autouse=True)
def v1(run_with_both_engines):
    # Simple autouse wrapper to run both engines for each test
    # This can be promoted up to conftest.py to run for every
    # test in a package
    pass


17
def test_empty_prompt():
18
19
20
    llm = LLM(model="s3://vllm-ci-model-weights/gpt2",
              load_format=LoadFormat.RUNAI_STREAMER,
              enforce_eager=True)
21
22
    with pytest.raises(ValueError, match='Prompt cannot be empty'):
        llm.generate([""])
23
24


Joe Runde's avatar
Joe Runde committed
25
@pytest.mark.skip_v1
26
def test_out_of_vocab_token():
27
28
29
    llm = LLM(model="s3://vllm-ci-model-weights/gpt2",
              load_format=LoadFormat.RUNAI_STREAMER,
              enforce_eager=True)
30
31
    with pytest.raises(ValueError, match='out of vocabulary'):
        llm.generate({"prompt_token_ids": [999999]})