From 5c6746f3ca00891a9096f8123218ca3242b794dd Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <oholworthy@nvidia.com>
Date: Tue, 15 Nov 2022 20:32:40 +0000
Subject: [PATCH] Update random sampling in tests to impove test reliability
 (#228)

---
 tests/unit/systems/dag/ops/test_ops.py   | 5 +++--
 tests/unit/systems/dag/test_ensemble.py  | 4 +++-
 tests/unit/systems/dag/test_executors.py | 7 ++++---
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/unit/systems/dag/ops/test_ops.py b/tests/unit/systems/dag/ops/test_ops.py
index ff0d8ce5c..d29255283 100644
--- a/tests/unit/systems/dag/ops/test_ops.py
+++ b/tests/unit/systems/dag/ops/test_ops.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import random
 from distutils.spawn import find_executable
 
 import numpy as np
@@ -38,7 +39,7 @@ def test_softmax_sampling(tmpdir):
     )
 
     combined_features = {
-        "movie_ids": np.random.randint(0, 10000, 100).astype(np.int32),
+        "movie_ids": np.array(random.sample(range(10000), 100), dtype=np.int32),
         "output_1": np.random.random(100).astype(np.float32),
     }
 
@@ -65,7 +66,7 @@ def test_filter_candidates_with_triton(tmpdir):
         ]
     )
 
-    candidate_ids = np.random.randint(1, 100000, 100).astype(np.int32)
+    candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32)
     movie_ids_1 = np.zeros(100, dtype=np.int32)
     movie_ids_1[:20] = np.unique(candidate_ids)[:20]
 
diff --git a/tests/unit/systems/dag/test_ensemble.py b/tests/unit/systems/dag/test_ensemble.py
index 45f1a6e3a..05a91e084 100644
--- a/tests/unit/systems/dag/test_ensemble.py
+++ b/tests/unit/systems/dag/test_ensemble.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import random
+
 import numpy as np
 
 from merlin.dag.executors import LocalExecutor
@@ -30,7 +32,7 @@ def test_ensemble_save_load(tmpdir):
         ]
     )
 
-    candidate_ids = np.random.randint(1, 100000, 100).astype(np.int32)
+    candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32)
     movie_ids_1 = np.zeros(100, dtype=np.int32)
     movie_ids_1[:20] = np.unique(candidate_ids)[:20]
 
diff --git a/tests/unit/systems/dag/test_executors.py b/tests/unit/systems/dag/test_executors.py
index 8c9f9f9fa..a4409d21f 100644
--- a/tests/unit/systems/dag/test_executors.py
+++ b/tests/unit/systems/dag/test_executors.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import random
 from distutils.spawn import find_executable
 
 import numpy as np
@@ -38,7 +39,7 @@ def test_run_dag_on_dictarray_with_local_executor():
         ]
     )
 
-    candidate_ids = np.random.randint(1, 100000, 100).astype(np.int32)
+    candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32)
     movie_ids_1 = np.zeros(100, dtype=np.int32)
     movie_ids_1[:20] = np.unique(candidate_ids)[:20]
 
@@ -72,7 +73,7 @@ def test_run_dag_on_dataframe_with_local_executor():
         ]
     )
 
-    candidate_ids = np.random.randint(1, 100000, 100).astype(np.int32)
+    candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32)
     movie_ids_1 = np.zeros(100, dtype=np.int32)
     movie_ids_1[:20] = np.unique(candidate_ids)[:20]
 
@@ -105,7 +106,7 @@ def test_run_dag_on_dataframe_with_dask_executor():
         ]
     )
 
-    candidate_ids = np.random.randint(1, 100000, 100).astype(np.int32)
+    candidate_ids = np.array(random.sample(range(100000), 100), dtype=np.int32)
     movie_ids_1 = np.zeros(100, dtype=np.int32)
     movie_ids_1[:20] = np.unique(candidate_ids)[:20]