From 8ccd31fe4952a77f63e506b740c1b32e56d1db60 Mon Sep 17 00:00:00 2001
From: SomberNight <somber.night@protonmail.com>
Date: Thu, 13 Mar 2025 18:30:53 +0000
Subject: [PATCH] wallet: set_frozen_state_of_coins to handle freeze=None

Internally whether a coin is frozen is tri-state:
- forced-True, set by the user
- forced-False, set by the user
- unset/default: is_frozen_coin() can decide whether the coin should be frozen

This patch lets set_frozen_state_of_coins() undo a previous explicit setting of True/False,
by calling it with a value of None.
Note: there is still no way in the GUI to undo an explicit setting of True/False.
---
 electrum/wallet.py            | 14 +++++++++---
 tests/test_wallet_vertical.py | 41 +++++++++++++++++++++++++++--------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/electrum/wallet.py b/electrum/wallet.py
index 9b72948c2..b6e026dab 100644
--- a/electrum/wallet.py
+++ b/electrum/wallet.py
@@ -2024,16 +2024,24 @@ class Abstract_Wallet(ABC, Logger, EventListener):
     def set_frozen_state_of_coins(
         self,
         utxos: Iterable[str],
-        freeze: bool,
+        freeze: Optional[bool],  # tri-state
         *,
         write_to_disk: bool = True,
     ) -> None:
-        """Set frozen state of the utxos to FREEZE, True or False"""
+        """Set frozen state of the utxos to `freeze`, True or False (or None).
+        A value of True/False means the user explicitly set if the coin should be frozen.
+        In contrast, None is the default "unset" state. If unset, is_frozen_coin()
+        can decide whether a coin should be frozen.
+        """
         # basic sanity check that input is not garbage: (see if raises)
         [TxOutpoint.from_str(utxo) for utxo in utxos]
+        assert freeze in (None, False, True), f"{freeze=!r}"
         with self._freeze_lock:
             for utxo in utxos:
-                self._frozen_coins[utxo] = bool(freeze)
+                if freeze is None:
+                    self._frozen_coins.pop(utxo, None)
+                else:
+                    self._frozen_coins[utxo] = bool(freeze)
         util.trigger_callback('status')
         if write_to_disk:
             self.save_db()
diff --git a/tests/test_wallet_vertical.py b/tests/test_wallet_vertical.py
index 028b7525c..f2411abfb 100644
--- a/tests/test_wallet_vertical.py
+++ b/tests/test_wallet_vertical.py
@@ -2895,16 +2895,39 @@ class TestWalletSending(ElectrumTestCase):
         self.assertEqual(
             {'52e669a20a26c8b3df5b41e5e6309b18bcde8e1ad7ea17a18f63b6dc6c8becc0:1'},
             {txi.prevout.to_str() for txi in wallet.get_spendable_coins(["tb1q6n99dl96mx8mfh90m3tn5awk5mllkzdh25dw7z"])})
+
+        utxo1 = "c36a6e1cd54df108e69574f70bc9b88dc13beddc70cfad9feb7f8f6593255d4a:1"
+        utxo2 = "52e669a20a26c8b3df5b41e5e6309b18bcde8e1ad7ea17a18f63b6dc6c8becc0:1"
+
         # test freezing an address
-        wallet.set_frozen_state_of_addresses(["tb1q6n99dl96mx8mfh90m3tn5awk5mllkzdh25dw7z"], freeze=True)
-        self.assertEqual(
-            {'c36a6e1cd54df108e69574f70bc9b88dc13beddc70cfad9feb7f8f6593255d4a:1'},
-            {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
-        wallet.set_frozen_state_of_addresses(["tb1q6n99dl96mx8mfh90m3tn5awk5mllkzdh25dw7z"], freeze=False)
-        self.assertEqual(
-            {'c36a6e1cd54df108e69574f70bc9b88dc13beddc70cfad9feb7f8f6593255d4a:1',
-             '52e669a20a26c8b3df5b41e5e6309b18bcde8e1ad7ea17a18f63b6dc6c8becc0:1'},
-            {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
+        with self.subTest(msg="freeze_address"):
+            wallet.set_frozen_state_of_addresses(["tb1q6n99dl96mx8mfh90m3tn5awk5mllkzdh25dw7z"], freeze=True)
+            self.assertEqual(
+                {utxo1},
+                {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
+            wallet.set_frozen_state_of_addresses(["tb1q6n99dl96mx8mfh90m3tn5awk5mllkzdh25dw7z"], freeze=False)
+            self.assertEqual(
+                {utxo1, utxo2},
+                {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
+
+        # test freezing a utxo
+        with self.subTest(msg="freeze_coin"):
+            self.assertTrue(utxo1 not in wallet._frozen_coins)
+
+            wallet.set_frozen_state_of_coins([utxo1], freeze=True)
+            self.assertEqual(wallet._frozen_coins.get(utxo1), True)
+            self.assertEqual(
+                {utxo2},
+                {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
+
+            wallet.set_frozen_state_of_coins([utxo1], freeze=False)
+            self.assertEqual(wallet._frozen_coins.get(utxo1), False)
+            self.assertEqual(
+                {utxo1, utxo2},
+                {txi.prevout.to_str() for txi in wallet.get_spendable_coins()})
+
+            wallet.set_frozen_state_of_coins([utxo1], freeze=None)
+            self.assertTrue(utxo1 not in wallet._frozen_coins)
 
     @mock.patch.object(wallet.Abstract_Wallet, 'save_db')
     async def test_export_psbt_with_xpubs__multisig(self, mock_save_db):