Skip to content

Commit ba46b92

Browse files
authored
Add default tag to functional tests (#166)
1 parent 5f4a3e4 commit ba46b92

File tree

1 file changed

+17
-2
lines changed

1 file changed

+17
-2
lines changed

test/functional/conftest.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import boto3
1717
import pytest
1818
from sagemaker import Session
19+
from sagemaker.tensorflow import TensorFlow
1920

2021
logger = logging.getLogger(__name__)
2122
logging.getLogger('boto').setLevel(logging.INFO)
@@ -31,6 +32,8 @@ def pytest_addoption(parser):
3132
parser.addoption('--instance-type')
3233
parser.addoption('--accelerator-type', default=None)
3334
parser.addoption('--region', default='us-west-2')
35+
parser.addoption('--framework-version', default=TensorFlow.LATEST_VERSION)
36+
parser.addoption('--processor', default='cpu', choices=['gpu', 'cpu'])
3437
parser.addoption('--tag')
3538

3639

@@ -60,8 +63,20 @@ def region(request):
6063

6164

6265
@pytest.fixture(scope='session')
63-
def tag(request):
64-
return request.config.getoption('--tag')
66+
def framework_version(request):
67+
return request.config.getoption('--framework-version')
68+
69+
70+
@pytest.fixture(scope='session')
71+
def processor(request):
72+
return request.config.getoption('--processor')
73+
74+
75+
@pytest.fixture(scope='session')
76+
def tag(request, framework_version, processor):
77+
provided_tag = request.config.getoption('--tag')
78+
default_tag = '{}-{}-py2'.format(framework_version, processor)
79+
return provided_tag if provided_tag is not None else default_tag
6580

6681

6782
@pytest.fixture(scope='session')

0 commit comments

Comments
 (0)