Merge pull request #79 from yzhu0/addRemoveTCSummary

fix: able to add and remove trail component summary

Author: yzhu0

src/smexperiments/trial.py

@@ -202,14 +202,16 @@ def add_trial_component(self, tc):
         A trial component may belong to many trials and a trial may have many trial components.
 
         Args:
-           tc: (tracker.Tracker|trial_component.TrialComponent|str) The trial component to add. Can be
-           one of a Tracker instance, a TrialComponent instance, or a string containing the name of
+           tc: (tracker.Tracker|trial_component.TrialComponent|api_types.TrialComponentSummary|str) The trial component
+           to add. Can be one of a Tracker instance, a TrialComponent instance, or a string containing the name of
            the trial component to add.
         """
         if isinstance(tc, tracker.Tracker):
             trial_component_name = tc.trial_component.trial_component_name
         elif isinstance(tc, trial_component.TrialComponent):
             trial_component_name = tc.trial_component_name
+        elif isinstance(tc, api_types.TrialComponentSummary):
+            trial_component_name = tc.trial_component_name
         else:
             trial_component_name = str(tc)
         self.sagemaker_boto_client.associate_trial_component(
@@ -220,14 +222,16 @@ def remove_trial_component(self, tc):
         """Remove the specified trial component from this trial.
 
         Args:
-            tc: (tracker.Tracker|trial_component.TrialComponent|str) The trial component to remove. Can be
-            one of a Tracker instance, a TrialComponent instance, or a string containing the name of
-            the trial component to remove.
+            tc: (tracker.Tracker|trial_component.TrialComponent|api_types.TrialComponentSummary|str) The trial
+            component to remove. Can be one of a Tracker instance, a TrialComponent instance, or a string containing
+            the name of the trial component to remove.
         """
         if isinstance(tc, tracker.Tracker):
             trial_component_name = tc.trial_component.trial_component_name
         elif isinstance(tc, trial_component.TrialComponent):
             trial_component_name = tc.trial_component_name
+        elif isinstance(tc, api_types.TrialComponentSummary):
+            trial_component_name = tc.trial_component_name
         else:
             trial_component_name = str(tc)
         self.sagemaker_boto_client.disassociate_trial_component(

tests/unit/test_trial.py

@@ -99,13 +99,31 @@ def test_add_trial_component(sagemaker_boto_client):
     )
 
 
+def test_add_trial_component_from_trial_component_summary(sagemaker_boto_client):
+    t = trial.Trial(sagemaker_boto_client)
+    t.trial_name = "bar"
+    tcs = api_types.TrialComponentSummary()
+    tcs.trial_component_name = "tcs-foo"
+    t.add_trial_component(tcs)
+    sagemaker_boto_client.associate_trial_component.assert_called_with(TrialName="bar", TrialComponentName="tcs-foo")
+
+
 def test_remove_trial_component(sagemaker_boto_client):
     t = trial.Trial(sagemaker_boto_client)
     t.trial_name = "bar"
     t.remove_trial_component("foo")
     sagemaker_boto_client.disassociate_trial_component.assert_called_with(TrialName="bar", TrialComponentName="foo")
 
 
+def test_remove_trial_component_from_trial_component_summary(sagemaker_boto_client):
+    t = trial.Trial(sagemaker_boto_client)
+    t.trial_name = "bar"
+    tcs = api_types.TrialComponentSummary()
+    tcs.trial_component_name = "tcs-foo"
+    t.remove_trial_component(tcs)
+    sagemaker_boto_client.disassociate_trial_component.assert_called_with(TrialName="bar", TrialComponentName="tcs-foo")
+
+
 def test_remove_trial_component_from_tracker(sagemaker_boto_client):
     t = trial.Trial(sagemaker_boto_client)
     t.trial_name = "bar"