16
16
import tempfile
17
17
import unittest
18
18
19
+ # https://github.com/coala/coala-bears/issues/2862
20
+ from unittest .mock import patch
21
+
19
22
import pytest
20
23
21
24
from smexperiments import _environment
@@ -36,7 +39,7 @@ def sagemaker_boto_client():
36
39
@pytest .fixture
37
40
def training_job_env ():
38
41
old_value = os .environ .get ("TRAINING_JOB_ARN" )
39
- os .environ ["TRAINING_JOB_ARN" ] = "arn:1234 "
42
+ os .environ ["TRAINING_JOB_ARN" ] = "arn:1234aBcDe "
40
43
yield os .environ
41
44
del os .environ ["TRAINING_JOB_ARN" ]
42
45
if old_value :
@@ -46,17 +49,17 @@ def training_job_env():
46
49
def test_processing_job_environment (tempdir ):
47
50
config_path = os .path .join (tempdir , "config.json" )
48
51
with open (config_path , "w" ) as f :
49
- f .write (json .dumps ({"ProcessingJobArn" : "arn:1234 " }))
52
+ f .write (json .dumps ({"ProcessingJobArn" : "arn:1234aBcDe " }))
50
53
environment = _environment .TrialComponentEnvironment .load (processing_job_config_path = config_path )
51
54
52
55
assert _environment .EnvironmentType .SageMakerProcessingJob == environment .environment_type
53
- assert "arn:1234 " == environment .source_arn
56
+ assert "arn:1234aBcDe " == environment .source_arn
54
57
55
58
56
59
def test_training_job_environment (training_job_env ):
57
60
environment = _environment .TrialComponentEnvironment .load ()
58
61
assert _environment .EnvironmentType .SageMakerTrainingJob == environment .environment_type
59
- assert "arn:1234 " == environment .source_arn
62
+ assert "arn:1234aBcDe " == environment .source_arn
60
63
61
64
62
65
def test_no_environment ():
@@ -70,9 +73,11 @@ def test_resolve_trial_component(training_job_env, sagemaker_boto_client):
70
73
}
71
74
sagemaker_boto_client .describe_trial_component .return_value = {"TrialComponentName" : trial_component_name }
72
75
environment = _environment .TrialComponentEnvironment .load ()
76
+
73
77
tc = environment .get_trial_component (sagemaker_boto_client )
74
78
75
79
assert trial_component_name == tc .trial_component_name
80
+ sagemaker_boto_client .list_trial_components .assert_called_with (SourceArn = "arn:1234abcde" )
76
81
sagemaker_boto_client .describe_trial_component .assert_called_with (TrialComponentName = trial_component_name )
77
82
78
83
0 commit comments