From 9d3b727dc1546e499f589092323d8a1020fc10e9 Mon Sep 17 00:00:00 2001
From: Dana Benson <danabens@amazon.com>
Date: Mon, 13 Jan 2020 21:07:29 -0800
Subject: [PATCH] add integ tests for list trial component by trial name and
 experiment name

---
 .../integ-jobs/test_track_from_processing_job.py  |  4 ++--
 tests/integ-jobs/test_track_from_training_job.py  |  3 ++-
 tests/integ/test_experiment.py                    |  1 +
 tests/integ/test_trial.py                         |  5 ++++-
 tests/integ/test_trial_component.py               | 15 +++++++++++++++
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/tests/integ-jobs/test_track_from_processing_job.py b/tests/integ-jobs/test_track_from_processing_job.py
index 4b8c6a2..c5d613c 100644
--- a/tests/integ-jobs/test_track_from_processing_job.py
+++ b/tests/integ-jobs/test_track_from_processing_job.py
@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 
 import sys
-
+import pytest
 from tests.helpers import *
 
 from smexperiments import trial_component, api_types
@@ -35,7 +35,7 @@ def wait_for_job(job, sagemaker_client):
                 sys.stdout.flush()
                 time.sleep(30)
 
-
+@pytest.mark.skip(reason="to be only run manually, integ/canaries already cover this scenario")
 def test_track_from_processing_job(sagemaker_boto_client, processing_job_name):
     processing_job = sagemaker_boto_client.describe_processing_job(ProcessingJobName=processing_job_name)
 
diff --git a/tests/integ-jobs/test_track_from_training_job.py b/tests/integ-jobs/test_track_from_training_job.py
index 1f35d17..37937fc 100644
--- a/tests/integ-jobs/test_track_from_training_job.py
+++ b/tests/integ-jobs/test_track_from_training_job.py
@@ -13,6 +13,7 @@
 
 import sys
 import boto3
+import pytest
 
 from tests.helpers import *
 from smexperiments import trial_component
@@ -58,7 +59,7 @@ def wait_for_job(job, sagemaker_client):
                 sys.stdout.flush()
                 time.sleep(30)
 
-
+@pytest.mark.skip(reason="to be only run manually, integ/canaries already cover this scenario")
 def test_track_from_training_job(sagemaker_boto_client, training_job_name):
     tj = sagemaker_boto_client.describe_training_job(TrainingJobName=training_job_name)
     source_arn = tj['TrainingJobArn']
diff --git a/tests/integ/test_experiment.py b/tests/integ/test_experiment.py
index bede613..fb899d2 100644
--- a/tests/integ/test_experiment.py
+++ b/tests/integ/test_experiment.py
@@ -94,3 +94,4 @@ def test_list_trials(experiment_obj, trials):
     assert set(trial_names) == set([s.trial_name for s in experiment_obj.list_trials()])
     assert trial_names  # sanity test
 
+
diff --git a/tests/integ/test_trial.py b/tests/integ/test_trial.py
index eb75e40..97e5c6f 100644
--- a/tests/integ/test_trial.py
+++ b/tests/integ/test_trial.py
@@ -55,8 +55,11 @@ def test_list_sort(trials, sagemaker_boto_client):
 
 def test_add_remove_trial_component(trial_obj, trial_component_obj):
     trial_obj.add_trial_component(trial_component_obj)
+    trial_components = list(trial_obj.list_trial_components())
+    assert 1 == len(trial_components)
     trial_obj.remove_trial_component(trial_component_obj)
-
+    trial_components = list(trial_obj.list_trial_components())
+    assert 0 == len(trial_components)
 
 def test_save(trial_obj, sagemaker_boto_client):
     trial_obj.display_name = 'foo'
diff --git a/tests/integ/test_trial_component.py b/tests/integ/test_trial_component.py
index aa068cb..f4bcc23 100644
--- a/tests/integ/test_trial_component.py
+++ b/tests/integ/test_trial_component.py
@@ -77,3 +77,18 @@ def test_list_sort(trial_components, sagemaker_boto_client):
         trial_component_names_listed = trial_component_names_listed[::-1]
     assert trial_component_names == trial_component_names_listed
     assert trial_component_names # sanity test
+
+
+def test_list_trial_components_by_experiment(experiment_obj, trial_component_obj, sagemaker_boto_client):
+    trial_obj = experiment_obj.create_trial()
+    trial_obj.add_trial_component(trial_component_obj)
+    trial_components = list(trial_component.TrialComponent.list(
+                            sagemaker_boto_client=sagemaker_boto_client,
+                            experiment_name=experiment_obj.experiment_name))
+    assert 1 == len(trial_components)
+    trial_obj.remove_trial_component(trial_component_obj)
+    trial_components = list(trial_component.TrialComponent.list(
+        sagemaker_boto_client=sagemaker_boto_client,
+        experiment_name=experiment_obj.experiment_name))
+    assert 0 == len(trial_components)
+    trial_obj.delete()
\ No newline at end of file