From fd39956e66632f5a1c04e39a50e0062a86a7b15b Mon Sep 17 00:00:00 2001
From: Conrad Holt <conrad.holt@protonmail.com>
Date: Mon, 26 Feb 2024 19:33:14 -0600
Subject: [PATCH] Quantum chess: remove assert_this_or_that test util

---
 unitary/quantum_chess/quantum_board_test.py | 10 ++--------
 unitary/quantum_chess/test_utils.py         | 11 -----------
 2 files changed, 2 insertions(+), 19 deletions(-)

diff --git a/unitary/quantum_chess/quantum_board_test.py b/unitary/quantum_chess/quantum_board_test.py
index d40430a4..71ad799b 100644
--- a/unitary/quantum_chess/quantum_board_test.py
+++ b/unitary/quantum_chess/quantum_board_test.py
@@ -26,7 +26,6 @@
 from unitary.quantum_chess.test_utils import (
     assert_sample_distribution,
     assert_samples_in,
-    assert_this_or_that,
     assert_prob_about,
     assert_fifty_fifty,
 )
@@ -182,10 +181,7 @@ def test_split_move(move_type, board):
             move_variant=enums.MoveVariant.BASIC,
         )
     )
-    samples = b.sample(100)
-    assert_this_or_that(
-        samples, u.squares_to_bitboard(["a3"]), u.squares_to_bitboard(["c1"])
-    )
+    assert_samples_in(b, [u.squares_to_bitboard(["a3"]), u.squares_to_bitboard(["c1"])])
     probs = b.get_probability_distribution(5000)
     assert_fifty_fifty(probs, qb.square_to_bit("a3"))
     assert_fifty_fifty(probs, qb.square_to_bit("c1"))
@@ -200,9 +196,7 @@ def test_split_move(move_type, board):
     )
     assert b.do_move(m)
     samples = b.sample(100)
-    assert_this_or_that(
-        samples, u.squares_to_bitboard(["a3"]), u.squares_to_bitboard(["d1"])
-    )
+    assert_samples_in(b, [u.squares_to_bitboard(["a3"]), u.squares_to_bitboard(["d1"])])
     probs = b.get_probability_distribution(5000)
     assert_fifty_fifty(probs, u.square_to_bit("a3"))
     assert_fifty_fifty(probs, u.square_to_bit("d1"))
diff --git a/unitary/quantum_chess/test_utils.py b/unitary/quantum_chess/test_utils.py
index a5d71958..bbc187f3 100644
--- a/unitary/quantum_chess/test_utils.py
+++ b/unitary/quantum_chess/test_utils.py
@@ -67,17 +67,6 @@ def assert_sample_distribution(b, probability_map, p_significant=1e-6):
     ), f"Observed {observed} far from expected {expected} (p = {p})"
 
 
-def assert_this_or_that(samples, this, that):
-    """Asserts all the samples are either equal to this or that,
-    and that one of each exists in the samples.
-    """
-    assert any(sample == this for sample in samples)
-    assert any(sample == that for sample in samples)
-    assert all(sample == this or sample == that for sample in samples), print_samples(
-        samples
-    )
-
-
 def assert_prob_about(probs, that, expected, atol=0.04):
     """Checks that the probability is within atol of the expected value."""
     assert probs[that] > expected - atol