From 1be796418934ae7370cb0ed501877db59e738106 Mon Sep 17 00:00:00 2001
From: Fabian Jahr <fjahr@protonmail.com>
Date: Wed, 1 Jun 2022 01:11:50 +0200
Subject: [PATCH] test, wallet: Add mempool rescan test for import RPCs

---
 test/functional/wallet_import_rescan.py | 57 ++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/test/functional/wallet_import_rescan.py b/test/functional/wallet_import_rescan.py
index d9acc8cea5b..085ad51c799 100755
--- a/test/functional/wallet_import_rescan.py
+++ b/test/functional/wallet_import_rescan.py
@@ -87,6 +87,7 @@ class Variant(collections.namedtuple("Variant", "call data address_type rescan p
         assert_equal(len(txs), self.expected_txs)
 
         addresses = self.node.listreceivedbyaddress(minconf=0, include_watchonly=True, address_filter=self.address['address'])
+
         if self.expected_txs:
             assert_equal(len(addresses[0]["txids"]), self.expected_txs)
 
@@ -98,13 +99,18 @@ class Variant(collections.namedtuple("Variant", "call data address_type rescan p
             assert_equal(tx["category"], "receive")
             assert_equal(tx["label"], self.label)
             assert_equal(tx["txid"], txid)
-            assert_equal(tx["confirmations"], 1 + current_height - confirmation_height)
-            assert "trusted" not in tx
+
+            # If no confirmation height is given, the tx is still in the
+            # mempool.
+            confirmations = (1 + current_height - confirmation_height) if confirmation_height else 0
+            assert_equal(tx["confirmations"], confirmations)
+            if confirmations:
+                assert "trusted" not in tx
 
             address, = [ad for ad in addresses if txid in ad["txids"]]
             assert_equal(address["address"], self.address["address"])
             assert_equal(address["amount"], self.expected_balance)
-            assert_equal(address["confirmations"], 1 + current_height - confirmation_height)
+            assert_equal(address["confirmations"], confirmations)
             # Verify the transaction is correctly marked watchonly depending on
             # whether the transaction pays to an imported public key or
             # imported private key. The test setup ensures that transaction
@@ -162,11 +168,12 @@ class ImportRescanTest(BitcoinTestFramework):
         self.import_deterministic_coinbase_privkeys()
         self.stop_nodes()
 
-        self.start_nodes()
+        self.start_nodes(extra_args=[["-whitelist=noban@127.0.0.1"]] * self.num_nodes)
         for i in range(1, self.num_nodes):
             self.connect_nodes(i, 0)
 
     def run_test(self):
+
         # Create one transaction on node 0 with a unique amount for
         # each possible type of wallet import RPC.
         for i, variant in enumerate(IMPORT_VARIANTS):
@@ -207,7 +214,7 @@ class ImportRescanTest(BitcoinTestFramework):
                 variant.check()
 
         # Create new transactions sending to each address.
-        for i, variant in enumerate(IMPORT_VARIANTS):
+        for variant in IMPORT_VARIANTS:
             variant.sent_amount = get_rand_amount()
             variant.sent_txid = self.nodes[0].sendtoaddress(variant.address["address"], variant.sent_amount)
             self.generate(self.nodes[0], 1)  # Generate one block for each send
@@ -223,6 +230,46 @@ class ImportRescanTest(BitcoinTestFramework):
             variant.expected_txs += 1
             variant.check(variant.sent_txid, variant.sent_amount, variant.confirmation_height)
 
+        self.log.info('Test that the mempool is rescanned as well if the rescan parameter is set to true')
+
+        # The late timestamp and pruned variants are not necessary when testing mempool rescan
+        mempool_variants = [variant for variant in IMPORT_VARIANTS if variant.rescan != Rescan.late_timestamp and not variant.prune]
+        # No further blocks are mined so the timestamp will stay the same
+        timestamp = self.nodes[0].getblockheader(self.nodes[0].getbestblockhash())["time"]
+
+        # Create one transaction on node 0 with a unique amount for
+        # each possible type of wallet import RPC.
+        for i, variant in enumerate(mempool_variants):
+            variant.label = "mempool label {} {}".format(i, variant)
+            variant.address = self.nodes[1].getaddressinfo(self.nodes[1].getnewaddress(
+                label=variant.label,
+                address_type=variant.address_type.value,
+            ))
+            variant.key = self.nodes[1].dumpprivkey(variant.address["address"])
+            variant.initial_amount = get_rand_amount()
+            variant.initial_txid = self.nodes[0].sendtoaddress(variant.address["address"], variant.initial_amount)
+            variant.confirmation_height = 0
+            variant.timestamp = timestamp
+
+        assert_equal(len(self.nodes[0].getrawmempool()), len(mempool_variants))
+        self.sync_mempools()
+
+        # For each variation of wallet key import, invoke the import RPC and
+        # check the results from getbalance and listtransactions.
+        for variant in mempool_variants:
+            self.log.info('Run import for mempool variant {}'.format(variant))
+            expect_rescan = variant.rescan == Rescan.yes
+            variant.node = self.nodes[2 + IMPORT_NODES.index(ImportNode(variant.prune, expect_rescan))]
+            variant.do_import(variant.timestamp)
+            if expect_rescan:
+                variant.expected_balance = variant.initial_amount
+                variant.expected_txs = 1
+                variant.check(variant.initial_txid, variant.initial_amount)
+            else:
+                variant.expected_balance = 0
+                variant.expected_txs = 0
+                variant.check()
+
 
 if __name__ == "__main__":
     ImportRescanTest().main()