tests: add unittests for LNWallet just in time opening

diff --git a/tests/test_lnwallet.py b/tests/test_lnwallet.py
index 1b7637d..3fbc544 100644
--- a/tests/test_lnwallet.py
+++ b/tests/test_lnwallet.py
@@ -1,14 +1,22 @@
 import logging
 import os
+import asyncio
+from unittest import mock
+from decimal import Decimal
 
+from electrum.address_synchronizer import TX_HEIGHT_LOCAL
 import electrum.trampoline
 from . import ElectrumTestCase
 from .test_lnchannel import create_test_channels
 
-from electrum.lnutil import RECEIVED, MIN_FINAL_CLTV_DELTA_ACCEPTED, LnFeatures
-from electrum.lntransport import LNPeerAddr
+from electrum.lnutil import RECEIVED, MIN_FINAL_CLTV_DELTA_ACCEPTED, serialize_htlc_key, LnFeatures
 from electrum.logging import console_stderr_handler
+from electrum.lntransport import LNPeerAddr
 from electrum.invoices import LN_EXPIRY_NEVER, PR_UNPAID
+from electrum.lnpeer import Peer
+from electrum.lnchannel import Channel, ChannelState
+from electrum.lnonion import OnionPacket, OnionRoutingFailure
+from electrum.crypto import sha256
 
 
 class TestLNWallet(ElectrumTestCase):
@@ -144,3 +152,265 @@ class TestLNWallet(ElectrumTestCase):
         lnaddr3, _ = wallet.get_bolt11_invoice(payment_info=pi2, message='test', fallback_address=None)
         hint_node_ids3 = {route[0][0] for route in lnaddr3.get_routing_info('r')}
         self.assertEqual(hint_node_ids3, {trampoline_pubkey})
+
+    async def test_open_channel_just_in_time_success(self):
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = True
+
+        next_peer = mock.Mock(spec=Peer)
+        next_chan = mock.Mock(spec=Channel)
+        next_chan.get_scid_or_local_alias.return_value = bytes(8)
+
+        funding_tx = mock.Mock()
+        funding_tx.txid.return_value = os.urandom(32).hex()
+
+        wallet.open_channel_with_peer = mock.AsyncMock(return_value=(next_chan, funding_tx))
+        wallet.network.try_broadcasting = mock.AsyncMock(return_value=True)
+
+        preimage = os.urandom(32)
+        payment_hash = sha256(preimage)
+
+        htlc = mock.Mock()
+        htlc.htlc_id = 0
+        next_peer.send_htlc.return_value = htlc
+
+        task = asyncio.create_task(wallet.open_channel_just_in_time(
+            next_peer=next_peer,
+            next_amount_msat_htlc=1000000,
+            next_cltv_abs=500,
+            payment_hash=payment_hash,
+            next_onion=mock.Mock(spec=OnionPacket)
+        ))
+
+        await asyncio.sleep(0.1)
+        wallet.save_preimage(payment_hash, preimage)
+        htlc_key = await task
+        htlc_key_correct = serialize_htlc_key(next_chan.get_scid_or_local_alias(), htlc.htlc_id)
+        self.assertEqual(htlc_key, htlc_key_correct)
+
+        wallet.open_channel_with_peer.assert_called_once()
+        next_peer.send_htlc.assert_called_once()
+        wallet.network.try_broadcasting.assert_called()
+
+    async def test_open_channel_just_in_time_failure_channel_open(self):
+        """The channel opening failed on the LSP side because the client rejected the incoming channel"""
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = True
+        next_peer = mock.Mock(spec=Peer)
+        wallet.open_channel_with_peer = mock.AsyncMock(side_effect=Exception("peer rejected incoming channel"))
+        preimage = os.urandom(32)
+        wallet.save_preimage(sha256(preimage), preimage)
+        wallet._cleanup_failed_jit_channel = mock.AsyncMock()
+
+        with self.assertRaises(OnionRoutingFailure):
+            await wallet.open_channel_just_in_time(
+                next_peer=next_peer,
+                next_amount_msat_htlc=1000000,
+                next_cltv_abs=500,
+                payment_hash=sha256(preimage),
+                next_onion=mock.Mock(spec=OnionPacket)
+            )
+
+        self.assertIsNone(wallet.get_preimage(sha256(preimage)))
+        wallet._cleanup_failed_jit_channel.assert_not_called()
+
+    async def test_open_channel_just_in_time_failure_send_htlc(self):
+        """The LSP fails to forward the htlc to the client"""
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = True
+
+        next_peer = mock.Mock(spec=Peer)
+        chan = mock.Mock(spec=Channel)
+        funding_tx = mock.Mock()
+
+        wallet.open_channel_with_peer = mock.AsyncMock(return_value=(chan, funding_tx))
+        next_peer.send_htlc.side_effect = Exception("couldn't send htlc, peer disconnected")
+        preimage = os.urandom(32)
+        wallet.save_preimage(sha256(preimage), preimage)
+        wallet._cleanup_failed_jit_channel = mock.AsyncMock()
+
+        with self.assertRaises(OnionRoutingFailure):
+            await wallet.open_channel_just_in_time(
+                next_peer=next_peer,
+                next_amount_msat_htlc=1000000,
+                next_cltv_abs=500,
+                payment_hash=sha256(preimage),
+                next_onion=mock.Mock(spec=OnionPacket)
+            )
+
+        self.assertIsNone(wallet.get_preimage(sha256(preimage)))
+        wallet._cleanup_failed_jit_channel.assert_called_once_with(chan)
+
+    async def test_open_channel_just_in_time_failure_preimage_timeout(self):
+        """The client never releases the preimage"""
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = True
+
+        next_peer = mock.Mock(spec=Peer)
+        chan = mock.Mock(spec=Channel)
+        funding_tx = mock.Mock()
+
+        wallet.open_channel_with_peer = mock.AsyncMock(return_value=(chan, funding_tx))
+
+        htlc = mock.Mock()
+        next_peer.send_htlc.return_value = htlc
+
+        wallet._cleanup_failed_jit_channel = mock.AsyncMock()
+
+        with mock.patch('electrum.lnworker.LN_P2P_NETWORK_TIMEOUT', 0.01):
+            with self.assertRaises(OnionRoutingFailure):
+                await wallet.open_channel_just_in_time(
+                    next_peer=next_peer,
+                    next_amount_msat_htlc=1000000,
+                    next_cltv_abs=500,
+                    payment_hash=os.urandom(32),
+                    next_onion=mock.Mock(spec=OnionPacket)
+                )
+
+        wallet._cleanup_failed_jit_channel.assert_called_once_with(chan)
+
+    async def test_open_channel_just_in_time_failure_broadcast(self):
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = True
+
+        next_peer = mock.Mock(spec=Peer)
+        chan = mock.Mock(spec=Channel)
+
+        funding_tx = mock.Mock()
+
+        wallet.open_channel_with_peer = mock.AsyncMock(return_value=(chan, funding_tx))
+
+        preimage = os.urandom(32)
+        wallet.save_preimage(sha256(preimage), preimage)
+
+        wallet.network.try_broadcasting = mock.AsyncMock(return_value=False)
+        wallet.wallet.adb.get_tx_height = mock.Mock(return_value=mock.Mock(height=lambda: TX_HEIGHT_LOCAL))
+
+        wallet._cleanup_failed_jit_channel = mock.AsyncMock()
+
+        with mock.patch('electrum.lnworker.ZEROCONF_TIMEOUT', 0.01), \
+             mock.patch('electrum.lnworker.asyncio.sleep', new_callable=mock.AsyncMock):
+             with self.assertRaises(OnionRoutingFailure):
+                await wallet.open_channel_just_in_time(
+                    next_peer=next_peer,
+                    next_amount_msat_htlc=1000000,
+                    next_cltv_abs=500,
+                    payment_hash=sha256(preimage),
+                    next_onion=mock.Mock(spec=OnionPacket)
+                )
+
+        self.assertIsNone(wallet.get_preimage(sha256(preimage)))
+        wallet._cleanup_failed_jit_channel.assert_called_once_with(chan)
+
+    async def test_open_channel_just_in_time_config_disabled(self):
+        """open_channel_just_in_time rejects to open a channel if the config is disabled"""
+        wallet = self.lnwallet_anchors
+        wallet.config.ZEROCONF_MIN_OPENING_FEE = 0
+        wallet.config.OPEN_ZEROCONF_CHANNELS = False
+
+        with self.assertRaises(AssertionError):
+            await wallet.open_channel_just_in_time(
+                next_peer=mock.Mock(spec=Peer),
+                next_amount_msat_htlc=1000000,
+                next_cltv_abs=500,
+                payment_hash=os.urandom(32),
+                next_onion=mock.Mock(spec=OnionPacket)
+            )
+
+    async def test_cleanup_failed_jit_channel(self):
+        wallet = self.lnwallet_anchors
+
+        chan = mock.Mock(spec=Channel)
+        chan_id = os.urandom(32).hex()
+        chan.channel_id = chan_id
+        funding_txid = os.urandom(32).hex()
+        chan.funding_outpoint = mock.Mock()
+        chan.funding_outpoint.txid = funding_txid
+        chan.get_funding_height.return_value = None
+
+        # close_channel fails with exception
+        wallet.close_channel = mock.AsyncMock(side_effect=Exception("peer disconnected"))
+        wallet.remove_channel = mock.Mock()
+        wallet.lnwatcher = mock.Mock()
+        wallet.lnwatcher.adb = mock.Mock()
+        wallet.lnwatcher.adb.remove_transaction = mock.Mock()
+
+        await wallet._cleanup_failed_jit_channel(chan)
+
+        wallet.close_channel.assert_called_once_with(chan_id)
+        chan.set_state.assert_called_once_with(ChannelState.REDEEMED, force=True)
+        wallet.lnwatcher.adb.remove_transaction.assert_called_once_with(funding_txid)
+        wallet.remove_channel.assert_called_once_with(chan_id)
+
+    async def test_receive_requires_jit_channel(self):
+        wallet = self.lnwallet_anchors
+
+        with self.subTest(msg="cannot get jit channel"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=False)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(0))
+            self.assertFalse(wallet.receive_requires_jit_channel(1_000_000))
+
+        with self.subTest(msg="could get zeroconf channel but doesn't need one"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(2000))
+            self.assertFalse(wallet.receive_requires_jit_channel(1_000_000))
+
+        with self.subTest(msg="could get zeroconf channel and needs one"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(500))
+            self.assertTrue(wallet.receive_requires_jit_channel(1_000_000))
+
+        with self.subTest(msg="could get one but can receive exactly the requested amount"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(1000))
+            self.assertFalse(wallet.receive_requires_jit_channel(1_000_000))
+
+        with self.subTest(msg="0 amount invoice, could get channel but can receive something"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(1))
+            self.assertFalse(wallet.receive_requires_jit_channel(None))
+
+        with self.subTest(msg="0 amount invoice (None amount), cannot receive anything and can get channel"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(0))
+            self.assertTrue(wallet.receive_requires_jit_channel(None))
+
+        with self.subTest(msg="0 amount invoice (0 msat), cannot receive anything, could get channel"):
+            wallet.can_get_zeroconf_channel = mock.Mock(return_value=True)
+            wallet.num_sats_can_receive = mock.Mock(return_value=Decimal(0))
+            self.assertTrue(wallet.receive_requires_jit_channel(0))
+
+    async def test_can_get_zeroconf_channel(self):
+        wallet = self.lnwallet_anchors
+        valid_peer = "02" * 33 + "@localhost:9735"
+
+        with self.subTest(msg="disabled in config"):
+            wallet.config.OPEN_ZEROCONF_CHANNELS = False
+            wallet.config.ZEROCONF_TRUSTED_NODE = valid_peer
+            self.assertFalse(wallet.can_get_zeroconf_channel())
+
+        with self.subTest(msg="enabled, but no trusted node configured"):
+            wallet.config.OPEN_ZEROCONF_CHANNELS = True
+            wallet.config.ZEROCONF_TRUSTED_NODE = ''
+            self.assertFalse(wallet.can_get_zeroconf_channel())
+
+        with self.subTest(msg="enabled, invalid trusted node string"):
+            wallet.config.OPEN_ZEROCONF_CHANNELS = True
+            wallet.config.ZEROCONF_TRUSTED_NODE = "invalid_node_string"
+            self.assertFalse(wallet.can_get_zeroconf_channel())
+
+        with self.subTest(msg="enabled, valid trusted node, but not connected"):
+            wallet.config.OPEN_ZEROCONF_CHANNELS = True
+            wallet.config.ZEROCONF_TRUSTED_NODE = valid_peer
+            self.assertFalse(wallet.can_get_zeroconf_channel())
+
+        with self.subTest(msg="enabled, valid trusted node, and connected"):
+            wallet.lnpeermgr.get_peer_by_pubkey = mock.Mock(return_value=mock.Mock(spec=Peer))
+            wallet.config.OPEN_ZEROCONF_CHANNELS = True
+            wallet.config.ZEROCONF_TRUSTED_NODE = valid_peer
+            self.assertTrue(wallet.can_get_zeroconf_channel())