Fix exception handling in squeak_core payment subscription (#2204)

squeaknode/core/squeak_core.py

@@ -24,13 +24,11 @@ from typing import Callable
 from typing import Optional
 from typing import Tuple
 
-import grpc
 from bitcoin.core import CBlockHeader
 from squeak.core import CBaseSqueak
 from squeak.core import CSqueak
 
 from squeaknode.bitcoin.bitcoin_client import BitcoinClient
-from squeaknode.core.exception import InvoiceSubscriptionError
 from squeaknode.core.lightning_address import LightningAddressHostPort
 from squeaknode.core.offer import Offer
 from squeaknode.core.peer_address import PeerAddress
@@ -415,25 +413,20 @@ class SqueakCore:
 
         def get_payment_stream():
             # Yield the received payments.
-            try:
-                for invoice in invoice_stream.result_stream:
-                    if invoice.settled:
-                        payment_hash = invoice.r_hash
-                        settle_index = invoice.settle_index
-                        sent_offer = get_sent_offer_fn(payment_hash)
-                        if sent_offer is not None:
-                            yield ReceivedPayment(
-                                received_payment_id=None,
-                                created_time_ms=None,
-                                squeak_hash=sent_offer.squeak_hash,
-                                payment_hash=sent_offer.payment_hash,
-                                price_msat=sent_offer.price_msat,
-                                settle_index=settle_index,
-                                peer_address=sent_offer.peer_address,
-                            )
-            except grpc.RpcError as e:
-                if e.code() != grpc.StatusCode.CANCELLED:
-                    raise InvoiceSubscriptionError()
+            for invoice in invoice_stream.result_stream:
+                payment_hash = invoice.r_hash
+                settle_index = invoice.settle_index
+                sent_offer = get_sent_offer_fn(payment_hash)
+                if sent_offer is not None:
+                    yield ReceivedPayment(
+                        received_payment_id=None,
+                        created_time_ms=None,
+                        squeak_hash=sent_offer.squeak_hash,
+                        payment_hash=sent_offer.payment_hash,
+                        price_msat=sent_offer.price_msat,
+                        settle_index=settle_index,
+                        peer_address=sent_offer.peer_address,
+                    )
 
         return ReceivedPaymentsStream(
             cancel_fn=cancel_subscription,

squeaknode/lightning/lnd_lightning_client.py

@@ -27,6 +27,7 @@ import grpc
 
 from proto import lnd_pb2
 from proto import lnd_pb2_grpc
+from squeaknode.core.exception import InvoiceSubscriptionError
 from squeaknode.lightning.info import Info
 from squeaknode.lightning.invoice import Invoice
 from squeaknode.lightning.invoice_stream import InvoiceStream
@@ -165,7 +166,17 @@ class LNDLightningClient(LightningClient):
         subscribe_result = self.stub.SubscribeInvoices(
             subscribe_invoices_request,
         )
+
+        def get_invoice_stream():
+            try:
+                for invoice in subscribe_result:
+                    if invoice.settled:
+                        yield invoice
+            except grpc.RpcError as e:
+                if e.code() != grpc.StatusCode.CANCELLED:
+                    raise InvoiceSubscriptionError()
+
         return InvoiceStream(
             cancel=subscribe_result.cancel,
-            result_stream=iter(subscribe_result),
+            result_stream=get_invoice_stream(),
         )

tests/core/test_squeak_core.py

@@ -131,9 +131,10 @@ def failed_payment(payment_request):
 
 @pytest.fixture
 def invoice_stream(settled_invoice, unsettled_invoice):
+    # Length of invoices list is 1.
     invoices = [
         settled_invoice,
-        unsettled_invoice,
+        # unsettled_invoice,
     ]
     yield InvoiceStream(
         cancel=lambda: None,

tests/lightning/test_lnd_lightning_client.py

@@ -109,7 +109,7 @@ def rpc_invoice(preimage, payment_hash, payment_request, price_msat, creation_da
         r_hash=payment_hash,
         payment_request=payment_request,
         value_msat=price_msat,
-        settled=False,
+        settled=True,
         settle_index=0,
         creation_date=creation_date,
         expiry=expiry,
@@ -122,7 +122,7 @@ def invoice(payment_hash, payment_request, price_msat, creation_date, expiry):
         r_hash=payment_hash,
         payment_request=payment_request,
         value_msat=price_msat,
-        settled=False,
+        settled=True,
         settle_index=0,
         creation_date=creation_date,
         expiry=expiry,