From 9d8914572179c6721fd9bbfe8934f07c79f7d52a Mon Sep 17 00:00:00 2001 From: crawlingcub <86861129+crawlingcub@users.noreply.github.com> Date: Fri, 24 Sep 2021 13:12:06 -0500 Subject: [PATCH] Fixing flaky test --- tests/ampligraph/evaluation/test_protocol.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ampligraph/evaluation/test_protocol.py b/tests/ampligraph/evaluation/test_protocol.py index 0294b3fb..2581a17a 100644 --- a/tests/ampligraph/evaluation/test_protocol.py +++ b/tests/ampligraph/evaluation/test_protocol.py @@ -1079,7 +1079,7 @@ def test_select_best_model_ranking_random(): assert best_params['k'] in (2, 50) assert np.log(1.00001) <= best_params['optimizer_params']['lr'] <= np.log(100) assert len(experimental_history) == 10 - assert set(i["model_params"]["k"] for i in experimental_history) == {2, 50} + assert all(i["model_params"]["k"] in [2, 50] for i in experimental_history) assert np.all([np.log(1.00001) <= i["model_params"]["optimizer_params"]["lr"] <= np.log(100) for i in experimental_history]) assert len(set(frozenset(_flatten_nested_keys(i["model_params"]).items()) for i in experimental_history)) == 10