66import shutil
77
88import shortuuid
9- from pytest import Parser , Session , FixtureRequest , FixtureDef , Item , Config , CollectReport
9+ from pytest import (
10+ Parser ,
11+ Session ,
12+ FixtureRequest ,
13+ FixtureDef ,
14+ Item ,
15+ Config ,
16+ CollectReport ,
17+ )
1018from _pytest .terminal import TerminalReporter
1119from typing import Optional , Any
1220from pytest_testconfig import config as py_config
@@ -39,18 +47,26 @@ def pytest_addoption(parser: Parser) -> None:
3947
4048 # Buckets options
4149 buckets_group .addoption (
42- "--ci-s3-bucket-name" , default = os .environ .get ("CI_S3_BUCKET_NAME" ), help = "Ci S3 bucket name"
50+ "--ci-s3-bucket-name" ,
51+ default = os .environ .get ("CI_S3_BUCKET_NAME" ),
52+ help = "Ci S3 bucket name" ,
4353 )
4454 buckets_group .addoption (
45- "--ci-s3-bucket-region" , default = os .environ .get ("CI_S3_BUCKET_REGION" ), help = "Ci S3 bucket region"
55+ "--ci-s3-bucket-region" ,
56+ default = os .environ .get ("CI_S3_BUCKET_REGION" ),
57+ help = "Ci S3 bucket region" ,
4658 )
4759
4860 buckets_group .addoption (
49- "--ci-s3-bucket-endpoint" , default = os .environ .get ("CI_S3_BUCKET_ENDPOINT" ), help = "Ci S3 bucket endpoint"
61+ "--ci-s3-bucket-endpoint" ,
62+ default = os .environ .get ("CI_S3_BUCKET_ENDPOINT" ),
63+ help = "Ci S3 bucket endpoint" ,
5064 )
5165
5266 buckets_group .addoption (
53- "--models-s3-bucket-name" , default = os .environ .get ("MODELS_S3_BUCKET_NAME" ), help = "Models S3 bucket name"
67+ "--models-s3-bucket-name" ,
68+ default = os .environ .get ("MODELS_S3_BUCKET_NAME" ),
69+ help = "Models S3 bucket name" ,
5470 )
5571 buckets_group .addoption (
5672 "--models-s3-bucket-region" ,
@@ -91,6 +107,11 @@ def pytest_addoption(parser: Parser) -> None:
91107 action = "store_true" ,
92108 help = "Delete pre-upgrade resources; useful when debugging pre-upgrade tests" ,
93109 )
110+ upgrade_group .addoption (
111+ "--upgrade-deployment-modes" ,
112+ help = "Coma-separated str; specify inference service deployment modes tests to run in upgrade tests. "
113+ "If not set, all will be tested." ,
114+ )
94115
95116
96117def pytest_cmdline_main (config : Any ) -> None :
@@ -102,19 +123,45 @@ def pytest_collection_modifyitems(session: Session, config: Config, items: list[
102123 Pytest fixture to filter or re-order the items in-place.
103124
104125 Filters upgrade tests based on '--pre-upgrade' / '--post-upgrade' option and marker.
126+ If `--upgrade-deployment-modes` option is set, only tests with the specified deployment modes will be added.
105127 """
128+
129+ def _add_upgrade_test (_item : Item , _upgrade_deployment_modes : list [str ]) -> bool :
130+ """
131+ Add upgrade test to the list of tests to run.
132+
133+ Args:
134+ _item (Item): The test item.
135+ _upgrade_deployment_modes (list[str]): The deployment modes to test.
136+
137+ Returns:
138+ True if the test should be added, False otherwise.
139+
140+ """
141+ if not _upgrade_deployment_modes :
142+ return True
143+
144+ return any ([keyword for keyword in _item .keywords if keyword in _upgrade_deployment_modes ])
145+
106146 pre_upgrade_tests : list [Item ] = []
107147 post_upgrade_tests : list [Item ] = []
108148 non_upgrade_tests : list [Item ] = []
149+ upgrade_deployment_modes : list [str ] = []
109150
110151 run_pre_upgrade_tests : str | None = config .getoption (name = "pre_upgrade" )
111152 run_post_upgrade_tests : str | None = config .getoption (name = "post_upgrade" )
153+ if config_upgrade_deployment_modes := config .getoption (name = "upgrade_deployment_modes" ):
154+ upgrade_deployment_modes = config_upgrade_deployment_modes .split ("," )
112155
113156 for item in items :
114- if "pre_upgrade" in item .keywords :
157+ if "pre_upgrade" in item .keywords and _add_upgrade_test (
158+ _item = item , _upgrade_deployment_modes = upgrade_deployment_modes
159+ ):
115160 pre_upgrade_tests .append (item )
116161
117- elif "post_upgrade" in item .keywords :
162+ elif "post_upgrade" in item .keywords and _add_upgrade_test (
163+ _item = item , _upgrade_deployment_modes = upgrade_deployment_modes
164+ ):
118165 post_upgrade_tests .append (item )
119166
120167 else :
0 commit comments