From 46499c9a789c62e44f2d559e39f01da37872db5a Mon Sep 17 00:00:00 2001
From: Cameron Fairchild <cameron@opentensor.ai>
Date: Thu, 29 Sep 2022 11:51:41 -0400
Subject: [PATCH] add test for blacklisted stake

---
 .../unit_tests/bittensor_tests/test_neuron.py | 108 ++++++++++++++++--
 1 file changed, 96 insertions(+), 12 deletions(-)

diff --git a/tests/unit_tests/bittensor_tests/test_neuron.py b/tests/unit_tests/bittensor_tests/test_neuron.py
index 4d5195cb1f..d795dd80c0 100644
--- a/tests/unit_tests/bittensor_tests/test_neuron.py
+++ b/tests/unit_tests/bittensor_tests/test_neuron.py
@@ -1,5 +1,6 @@
 from atexit import register
 from types import SimpleNamespace
+import unittest
 from unittest.mock import MagicMock, patch
 from more_itertools import side_effect
 
@@ -77,9 +78,6 @@ def test_coreserver_reregister_flag_false_exit():
     mock_wallet = bittensor.wallet.mock()
     mock_wallet.config = config
 
-    class MockException(Exception):
-        pass
-
     def exit_early(*args, **kwargs):
         raise MockException('exit_early')
 
@@ -129,9 +127,6 @@ def test_coreserver_reregister_flag_true():
     mock_wallet = bittensor.wallet.mock()
     mock_wallet.config = config
 
-    class MockException(Exception):
-        pass
-
     def exit_early(*args, **kwargs):
         raise MockException('exit_early')
 
@@ -178,9 +173,6 @@ def test_corevalidator_reregister_flag_false_exit():
     mock_wallet = bittensor.wallet.mock()
     mock_wallet.config = config
 
-    class MockException(Exception):
-        pass
-
     def exit_early(*args, **kwargs):
         raise MockException('exit_early')
 
@@ -227,9 +219,6 @@ def test_corevalidator_reregister_flag_true():
     mock_wallet = bittensor.wallet.mock()
     mock_wallet.config = config
 
-    class MockException(Exception):
-        pass
-
     def exit_early(*args, **kwargs):
         raise MockException('exit_early')
 
@@ -260,5 +249,100 @@ def exit_early(*args, **kwargs):
         # Should try to register the neuron
         mock_register.assert_called_once()
 
+class MockException(Exception):
+    pass
+
+class TestBlacklist(unittest.TestCase):
+
+    @staticmethod
+    def construct_config():
+        defaults = bittensor.Config()
+        bittensor.subtensor.add_defaults( defaults )
+        bittensor.dendrite.add_defaults( defaults )
+        bittensor.axon.add_defaults( defaults )
+        bittensor.wallet.add_defaults( defaults )
+        bittensor.dataset.add_defaults( defaults )
+        bittensor.logging.add_defaults( defaults )
+        bittensor.wandb.add_defaults( defaults )
+        defaults.wandb.api_key = 'test'
+        defaults.neuron = bittensor.neurons.core_server.neuron.config()
+        defaults.neuron.learning_rate = 0.0001
+        defaults.neuron.momentum = 0.9
+        
+        return defaults
+    
+    def exit_early(self, *args, **kwargs):
+        raise MockException('exit_early')
+
+    def test_stake_blacklist(self):
+        import sys
+        sys.setrecursionlimit(100)
+
+        mock_hotkey = "0x0000000000000000000000000000000000000000"
+        mock_hotkey_1 = "0x0000000000000000000000000000000000000001"
+
+        mock_subtensor = MagicMock(
+            is_hotkey_registered=MagicMock(return_value=True),
+        )
+
+        mock_wallet = MagicMock(
+            reregister=MagicMock(),
+            is_registered=MagicMock(return_value=True),
+        )
+
+        mock_metagraph = MagicMock(
+            hotkeys=[
+                mock_hotkey,
+                mock_hotkey_1,
+            ],
+            S=[
+                torch.tensor(100), # stake for mock_hotkey, uid 0
+                torch.tensor(1001), # stake for mock_hotkey_1, uid 1
+            ]
+        )
+
+        mock_config = self.construct_config()
+        
+        mock_config.neuron.blacklist = bittensor.Config()
+        mock_config.neuron.blacklist.stake = 1000 # blacklist if stake is less than 1000
+
+        mock_model_config = bittensor.neurons.core_server.server.config()
+        mock_model_config.neuron = MagicMock(
+            disable_blacklist = False
+        )
+
+        mock_model = MagicMock(
+                            spec=bittensor.neurons.core_server.server,
+                            config=mock_model_config,
+                        )
+
+        with patch('bittensor.axon.__new__', side_effect=self.exit_early) as mock_new_axon:
+            with patch('bittensor.neurons.core_server.neuron.check_config', return_value=True):
+                with pytest.raises(MockException):
+                    bittensor.neurons.core_server.serve(
+                        config=mock_config,
+                        model=MagicMock(
+                            spec=bittensor.neurons.core_server.server,
+                            device="cpu",
+                            to=MagicMock(return_value=mock_model),
+                            config=mock_model_config,
+                        ),
+                        subtensor=mock_subtensor,
+                        wallet=mock_wallet,
+                        axon=None,
+                        metagraph=mock_metagraph
+                    )
+
+            # args, kwargs
+            _, kwargs = mock_new_axon.call_args
+            blacklist = kwargs['blacklist']
+
+            # Check that the blacklist rejects below min stake
+            assert blacklist(mock_hotkey, bittensor.proto.RequestType.FORWARD) == True
+
+            # Check that the blacklist accepts above min stake
+            assert blacklist(mock_hotkey_1, bittensor.proto.RequestType.FORWARD) == False
+
+
 if __name__ == '__main__':
     pass