From b5969ceaccd7f44920bfd282861642fd6f5f9d48 Mon Sep 17 00:00:00 2001 From: Dana Benson Date: Mon, 18 Oct 2021 12:26:39 -0700 Subject: [PATCH] retry downstream_trials test --- tests/integ/sagemaker/lineage/helpers.py | 14 +++++++++++++ .../integ/sagemaker/lineage/test_artifact.py | 21 ++++++++++++------- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py index 028379259b..757ee3bf6d 100644 --- a/tests/integ/sagemaker/lineage/helpers.py +++ b/tests/integ/sagemaker/lineage/helpers.py @@ -15,6 +15,7 @@ import uuid from datetime import datetime +import time def name(): @@ -30,3 +31,16 @@ def names(): ) for i in range(3) ] + + +def retry(callable, num_attempts=8): + assert num_attempts >= 1 + for i in range(num_attempts): + try: + return callable() + except Exception as ex: + if i == num_attempts - 1: + raise ex + print("Retrying", ex) + time.sleep(2 ** i) + assert False, "logic error in retry" diff --git a/tests/integ/sagemaker/lineage/test_artifact.py b/tests/integ/sagemaker/lineage/test_artifact.py index 9f47c07a98..4a0c6398b2 100644 --- a/tests/integ/sagemaker/lineage/test_artifact.py +++ b/tests/integ/sagemaker/lineage/test_artifact.py @@ -20,6 +20,7 @@ import pytest from sagemaker.lineage import artifact +from tests.integ.sagemaker.lineage.helpers import retry def test_create_delete(artifact_obj): @@ -103,14 +104,20 @@ def test_list_by_type(artifact_objs, sagemaker_session): def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_session): # allow trial components to index, 30 seconds max - for i in range(3): - time.sleep(10) - trials = trial_associated_artifact.downstream_trials(sagemaker_session=sagemaker_session) - if len(trials) > 0: - break + def validate(): + for i in range(3): + time.sleep(10) + trials = trial_associated_artifact.downstream_trials( + sagemaker_session=sagemaker_session + ) + logging.info(f"Found {len(trials)} downstream trials.") + if len(trials) > 0: + break + + assert len(trials) == 1 + assert trial_obj.trial_name in trials - assert len(trials) == 1 - assert trial_obj.trial_name in trials + retry(validate, num_attempts=3) @pytest.mark.timeout(30)