From a0efda573c2a12a4abe69d79b46bcaecce0614f4 Mon Sep 17 00:00:00 2001
From: derekpierre <derek.pierre@gmail.com>
Date: Tue, 14 Nov 2023 16:29:23 -0500
Subject: [PATCH] Update sampling tests to test exclusion logic.

---
 .../agents/test_taco_application_agent.py      | 17 +++++++++++++++++
 .../test_taco_child_application_agent.py       | 18 ++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/tests/acceptance/agents/test_taco_application_agent.py b/tests/acceptance/agents/test_taco_application_agent.py
index 31a0935e8..fe125d616 100644
--- a/tests/acceptance/agents/test_taco_application_agent.py
+++ b/tests/acceptance/agents/test_taco_application_agent.py
@@ -1,3 +1,5 @@
+import random
+
 import pytest
 from eth_utils import is_address
 
@@ -99,6 +101,7 @@ def test_get_swarm(taco_application_agent, staking_providers):
 
 @pytest.mark.usefixtures("staking_providers", "ursulas")
 def test_sample_staking_providers(taco_application_agent):
+    all_staking_providers = taco_application_agent.get_staking_providers()
     providers_population = taco_application_agent.get_staking_providers_population()
 
     with pytest.raises(taco_application_agent.NotEnoughStakingProviders):
@@ -116,13 +119,27 @@ def test_sample_staking_providers(taco_application_agent):
     ).draw(3)
     assert len(providers) == 3
     assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
+
+    # repeat for opposite blockchain light setting
     light = taco_application_agent.blockchain.is_light
     taco_application_agent.blockchain.is_light = not light
     providers = taco_application_agent.get_staking_provider_reservoir().draw(3)
     assert len(providers) == 3
     assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
     taco_application_agent.blockchain.is_light = light
 
+    # Use exclusion list
+    exclude_providers = random.choices(all_staking_providers, k=3)  # exclude 3 ursulas
+    providers = taco_application_agent.get_staking_provider_reservoir(
+        without=exclude_providers, pagination_size=1
+    ).draw(3)
+    assert len(providers) == 3
+    assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
+    assert len(set(providers).intersection(exclude_providers)) == 0
+
 
 def test_get_staking_provider_info(testerchain, taco_application_agent):
     staking_provider_account, operator_account, *other = testerchain.unassigned_accounts
diff --git a/tests/acceptance/agents/test_taco_child_application_agent.py b/tests/acceptance/agents/test_taco_child_application_agent.py
index ccf467440..fa9705722 100644
--- a/tests/acceptance/agents/test_taco_child_application_agent.py
+++ b/tests/acceptance/agents/test_taco_child_application_agent.py
@@ -1,3 +1,5 @@
+import random
+
 import pytest
 
 from nucypher.blockchain.eth.constants import NULL_ADDRESS
@@ -100,6 +102,7 @@ def test_get_staker_population(taco_child_application_agent, staking_providers):
 
 @pytest.mark.usefixtures("staking_providers", "ursulas")
 def test_sample_staking_providers(taco_child_application_agent):
+    all_staking_providers = taco_child_application_agent.get_staking_providers()
     providers_population = (
         taco_child_application_agent.get_staking_providers_population()
     )
@@ -112,6 +115,7 @@ def test_sample_staking_providers(taco_child_application_agent):
     providers = taco_child_application_agent.get_staking_provider_reservoir().draw(3)
     assert len(providers) == 3  # Three...
     assert len(set(providers)) == 3  # ...unique addresses
+    assert len(set(providers).intersection(all_staking_providers)) == 3
 
     # Same but with pagination
     providers = taco_child_application_agent.get_staking_provider_reservoir(
@@ -119,9 +123,23 @@ def test_sample_staking_providers(taco_child_application_agent):
     ).draw(3)
     assert len(providers) == 3
     assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
+
+    # repeat for opposite blockchain light setting
     light = taco_child_application_agent.blockchain.is_light
     taco_child_application_agent.blockchain.is_light = not light
     providers = taco_child_application_agent.get_staking_provider_reservoir().draw(3)
     assert len(providers) == 3
     assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
     taco_child_application_agent.blockchain.is_light = light
+
+    # Use exclusion list
+    exclude_providers = random.choices(all_staking_providers, k=2)  # exclude 2 ursulas
+    providers = taco_child_application_agent.get_staking_provider_reservoir(
+        without=exclude_providers, pagination_size=1
+    ).draw(3)
+    assert len(providers) == 3
+    assert len(set(providers)) == 3
+    assert len(set(providers).intersection(all_staking_providers)) == 3
+    assert len(set(providers).intersection(exclude_providers)) == 0