Fix tests and error messages.

SFDO-Tooling · Jun 3, 2022 · 98d9b04 · 98d9b04
1 parent ef63376
commit 98d9b04
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
diff --git a/snowfakery/row_history.py b/snowfakery/row_history.py
@@ -184,9 +184,17 @@ def __init__(
             self.random_func = randint
 
     def next(self) -> T.Union[ObjectReference, ObjectRow]:
-        return self.row_history.random_row_reference(
-            self.to, self.scope, self.random_func
-        )
+        try:
+            return self.row_history.random_row_reference(
+                self.to, self.scope, self.random_func
+            )
+        except StopIteration as e:
+            if self.random_func == self.unique_random:
+                raise exc.DataGenError(
+                    f"Cannot find an unused `{self.to}`` to link to"
+                ) from e
+            else:  # pragma no cover
+                raise e
 
     def unique_random(self, a, b):
         """Goal: use an Uniquifying RNG until all of its values have been

diff --git a/tests/test_references.py b/tests/test_references.py
@@ -462,7 +462,25 @@ def test_random_reference_unique(self, generated_rows):
                 to: A               #10
                 unique: true        #11
     """
-        generate(StringIO(yaml), reps=5)
+        generate(StringIO(yaml), stopping_criteria=StoppingCriteria("B", 25))
+        Bs = generated_rows.table_values("B")
+        # check that A_refs are unique
+        assert len(Bs) == len(set(b["A_ref"] for b in Bs)) == 25
+
+    def test_random_reference_unique__nickname(self, generated_rows):
+        yaml = """                  #1
+      - object: A                   #2
+        nickname: nicky             #3
+        count: 5                   #4
+      - object: B                   #5
+        count: 5                    #6
+        fields:                     #7
+            A_ref:                  #8
+              random_reference:     #9
+                to: nicky           #10
+                unique: true        #11
+    """
+        generate(StringIO(yaml), stopping_criteria=StoppingCriteria("B", 25))
         Bs = generated_rows.table_values("B")
         # check that A_refs are unique
         assert len(Bs) == len(set(b["A_ref"] for b in Bs)) == 25
@@ -479,8 +497,8 @@ def test_random_reference_unique__mismatch_count(self, generated_rows):
                 to: A               #10
                 unique: true        #11
     """
-        generate(StringIO(yaml))
-        assert 0, generated_rows.mock_calls
+        with pytest.raises(DataGenError, match="Cannot find an unused `A`"):
+            generate(StringIO(yaml))
 
     def test_random_reference_unique__with_continuation(
         self, generate_data_with_continuation, generated_rows