19
19
from botocore .credentials import Credentials
20
20
21
21
import base64
22
+ import logging
22
23
import json
23
24
import os
24
25
import subprocess
@@ -332,12 +333,14 @@ def test_check_output():
332
333
@patch ("sagemaker.local.data.get_data_source_instance" )
333
334
@patch ("subprocess.Popen" )
334
335
def test_train (
335
- popen , get_data_source_instance , retrieve_artifacts , cleanup , tmpdir , sagemaker_session
336
+ popen , get_data_source_instance , retrieve_artifacts , cleanup , tmpdir , sagemaker_session , caplog
336
337
):
337
338
data_source = Mock ()
338
339
data_source .get_root_dir .return_value = "foo"
339
340
get_data_source_instance .return_value = data_source
340
341
342
+ caplog .set_level (logging .INFO )
343
+
341
344
directories = [str (tmpdir .mkdir ("container-root" )), str (tmpdir .mkdir ("data" ))]
342
345
with patch (
343
346
"sagemaker.local.image._SageMakerContainer._create_tmp_folder" , side_effect = directories
@@ -388,6 +391,7 @@ def test_train(
388
391
389
392
retrieve_artifacts .assert_called_once ()
390
393
cleanup .assert_called_once ()
394
+ assert "[Masked]" in caplog .text
391
395
392
396
393
397
@patch ("sagemaker.local.local_session.LocalSession" , Mock ())
@@ -579,7 +583,8 @@ def test_container_does_not_enable_nvidia_docker_for_cpu_containers(sagemaker_se
579
583
@patch ("sagemaker.local.image._SageMakerContainer._prepare_serving_volumes" , Mock (return_value = []))
580
584
@patch ("shutil.copy" , Mock ())
581
585
@patch ("shutil.copytree" , Mock ())
582
- def test_serve (tmpdir , sagemaker_session ):
586
+ def test_serve (tmpdir , sagemaker_session , caplog ):
587
+ caplog .set_level (logging .INFO )
583
588
with patch (
584
589
"sagemaker.local.image._SageMakerContainer._create_tmp_folder" ,
585
590
return_value = str (tmpdir .mkdir ("container-root" )),
@@ -601,6 +606,7 @@ def test_serve(tmpdir, sagemaker_session):
601
606
for h in sagemaker_container .hosts :
602
607
assert config ["services" ][h ]["image" ] == image
603
608
assert config ["services" ][h ]["command" ] == "serve"
609
+ assert "[Masked]" in caplog .text
604
610
605
611
606
612
@patch ("sagemaker.local.image._HostingContainer.run" , Mock ())
@@ -765,7 +771,7 @@ def test_ecr_login_needed(check_output):
765
771
"docker login -u AWS -p %s https://520713654638.dkr.ecr.us-east-1.amazonaws.com" % token
766
772
)
767
773
768
- check_output .assert_called_with (expected_command .split ())
774
+ check_output .assert_called_with (expected_command .split (), stderr = subprocess . DEVNULL )
769
775
session_mock .client ("ecr" ).get_authorization_token .assert_called_with (
770
776
registryIds = ["520713654638" ]
771
777
)
0 commit comments