Fix region-handling for integ test sagemaker_session (#106)

* Fix region-handling for integ test sagemaker_session * Fix flake8 error
aws · Mar 23, 2018 · 9d33e6b · 9d33e6b
1 parent 447197e
commit 9d33e6b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -46,9 +46,10 @@ def boto_config(request):
 
 @pytest.fixture(scope='session')
 def sagemaker_session(sagemaker_client_config, sagemaker_runtime_config, boto_config):
-    sagemaker_client = boto3.client('sagemaker', **sagemaker_client_config) if sagemaker_client_config else None
-    runtime_client = boto3.client('sagemaker-runtime', **sagemaker_runtime_config) if sagemaker_runtime_config else None
     boto_session = boto3.Session(**boto_config) if boto_config else boto3.Session(region_name=DEFAULT_REGION)
+    sagemaker_client = boto_session.client('sagemaker', **sagemaker_client_config) if sagemaker_client_config else None
+    runtime_client = (boto_session.client('sagemaker-runtime', **sagemaker_runtime_config) if sagemaker_runtime_config
+                      else None)
 
     return Session(boto_session=boto_session,
                    sagemaker_client=sagemaker_client,