onion_message: move encrypt_onionmsg_tlv_hops_data() to lnonion as encrypt_hops_recipient_data() and add support payloads other than onionmsg_tlv

diff --git a/electrum/lnonion.py b/electrum/lnonion.py
index 06b500f..3117eb9 100644
--- a/electrum/lnonion.py
+++ b/electrum/lnonion.py
@@ -281,6 +281,29 @@ def decrypt_onionmsg_data_tlv(*, shared_secret: bytes, encrypted_recipient_data:
     return recipient_data
 
 
+def encrypt_hops_recipient_data(
+        tlv_stream_name: str,
+        hops_data: Sequence[OnionHopsDataSingle],
+        hop_shared_secrets: Sequence[bytes]
+) -> None:
+    """encrypt unencrypted encrypted_recipient_data for hops with blind_fields.
+
+       NOTE: contents of payload.encrypted_recipient_data is slightly different for 'payload'
+       vs 'oniomsg_tlv' tlv_stream_names, so we map to the correct key here based on tlv_stream_name.
+       We can also change onion_wire.csv to use the same key, but as we import that from specs it might
+       regress in the future, so I rather make it explicit in code here.
+    """
+    # key naming payload TLV vs onionmsg_tlv TLV
+    erd_key = 'encrypted_recipient_data' if tlv_stream_name == 'onionmsg_tlv' else 'encrypted_data'
+
+    num_hops = len(hops_data)
+    for i in range(num_hops):
+        if hops_data[i].tlv_stream_name == tlv_stream_name and 'encrypted_recipient_data' not in hops_data[i].payload:
+            # construct encrypted_recipient_data from blind_fields
+            encrypted_recipient_data = encrypt_onionmsg_data_tlv(shared_secret=hop_shared_secrets[i], **hops_data[i].blind_fields)
+            hops_data[i].payload['encrypted_recipient_data'] = {erd_key: encrypted_recipient_data}
+
+
 def calc_hops_data_for_payment(
         route: 'LNPaymentRoute',
         amount_msat: int,  # that final recipient receives
diff --git a/electrum/onion_message.py b/electrum/onion_message.py
index 29eef92..fcb1c14 100644
--- a/electrum/onion_message.py
+++ b/electrum/onion_message.py
@@ -41,7 +41,7 @@ from electrum.crypto import sha256, get_ecdh
 from electrum.lnmsg import OnionWireSerializer
 from electrum.lnonion import (get_bolt04_onion_key, OnionPacket, process_onion_packet,
                               OnionHopsDataSingle, decrypt_onionmsg_data_tlv, encrypt_onionmsg_data_tlv,
-                              get_shared_secrets_along_route, new_onion_packet)
+                              get_shared_secrets_along_route, new_onion_packet, encrypt_hops_recipient_data)
 from electrum.lnutil import LnFeatures
 from electrum.util import OldTaskGroup, log_exceptions
 
@@ -140,21 +140,6 @@ def is_onion_message_node(node_id: bytes, node_info: Optional['NodeInfo']) -> bo
     return LnFeatures(node_info.features).supports(LnFeatures.OPTION_ONION_MESSAGE_OPT)
 
 
-def encrypt_onionmsg_tlv_hops_data(
-        hops_data: List[OnionHopsDataSingle],
-        hop_shared_secrets: Sequence[bytes]
-) -> None:
-    """encrypt unencrypted onionmsg_tlv.encrypted_recipient_data for hops with blind_fields"""
-    num_hops = len(hops_data)
-    for i in range(num_hops):
-        if hops_data[i].tlv_stream_name == 'onionmsg_tlv' and 'encrypted_recipient_data' not in hops_data[i].payload:
-            # construct encrypted_recipient_data from blind_fields
-            encrypted_recipient_data = encrypt_onionmsg_data_tlv(shared_secret=hop_shared_secrets[i], **hops_data[i].blind_fields)
-            new_payload = dict(hops_data[i].payload)
-            new_payload['encrypted_recipient_data'] = {'encrypted_recipient_data': encrypted_recipient_data}
-            hops_data[i] = dataclasses.replace(hops_data[i], payload=new_payload)
-
-
 def create_onion_message_route_to(lnwallet: 'LNWallet', node_id: bytes) -> Sequence[PathEdge]:
     """Constructs a route to the destination node_id, first by starting with peers with existing channels,
        and if no route found, opening a direct peer connection if node_id is found with an address in
@@ -325,7 +310,7 @@ def send_onion_message_to(
 
             payment_path_pubkeys = blinded_node_ids + blinded_path_blinded_ids
             hop_shared_secrets, _ = get_shared_secrets_along_route(payment_path_pubkeys, session_key)
-            encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+            encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
             packet = new_onion_packet(payment_path_pubkeys, session_key, hops_data, onion_message=True)
             packet_b = packet.to_bytes()
 
@@ -365,7 +350,7 @@ def send_onion_message_to(
         payment_path_pubkeys = [edge.end_node for edge in path]
 
         hop_shared_secrets, blinded_node_ids = get_shared_secrets_along_route(payment_path_pubkeys, session_key)
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
         packet = new_onion_packet(blinded_node_ids, session_key, hops_data)
         packet_b = packet.to_bytes()
 
diff --git a/tests/test_onion_message.py b/tests/test_onion_message.py
index 2b57b77..4e381d6 100644
--- a/tests/test_onion_message.py
+++ b/tests/test_onion_message.py
@@ -13,15 +13,13 @@ from electrum_ecc import ECPrivkey
 from electrum import SimpleConfig
 from electrum.lnmsg import decode_msg, OnionWireSerializer
 from electrum.lnonion import (
-    OnionHopsDataSingle, OnionPacket,
-    process_onion_packet, get_bolt04_onion_key, encrypt_onionmsg_data_tlv,
-    get_shared_secrets_along_route, new_onion_packet, ONION_MESSAGE_LARGE_SIZE,
-    HOPS_DATA_SIZE, InvalidPayloadSize)
+    OnionHopsDataSingle, OnionPacket, process_onion_packet, get_bolt04_onion_key, encrypt_onionmsg_data_tlv,
+    get_shared_secrets_along_route, new_onion_packet, ONION_MESSAGE_LARGE_SIZE, HOPS_DATA_SIZE, InvalidPayloadSize,
+    encrypt_hops_recipient_data)
 from electrum.crypto import get_ecdh, privkey_to_pubkey
 from electrum.lnutil import LnFeatures, Keypair
 from electrum.onion_message import (
-    blinding_privkey, create_blinded_path, encrypt_onionmsg_tlv_hops_data,
-    OnionMessageManager, NoRouteFound, Timeout
+    blinding_privkey, create_blinded_path,OnionMessageManager, NoRouteFound, Timeout
 )
 from electrum.util import bfh, read_json_file, OldTaskGroup, get_asyncio_loop
 from electrum.logging import console_stderr_handler
@@ -104,7 +102,7 @@ class TestOnionMessage(ElectrumTestCase):
             )
         ]
 
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
         packet = new_onion_packet(blinded_node_ids, SESSION_KEY, hops_data, onion_message=True)
         self.assertEqual(packet.to_bytes(), ONION_MESSAGE_PACKET)
 
@@ -126,18 +124,19 @@ class TestOnionMessage(ElectrumTestCase):
                 ),
             ]
         hops_data = hops_data_for_message('short_message')  # fit in HOPS_DATA_SIZE
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
         packet = new_onion_packet(blinded_node_ids, SESSION_KEY, hops_data, onion_message=True)
         self.assertEqual(len(packet.to_bytes()), HOPS_DATA_SIZE + 66)
 
         hops_data = hops_data_for_message('A' * HOPS_DATA_SIZE)  # fit in ONION_MESSAGE_LARGE_SIZE
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
         packet = new_onion_packet(blinded_node_ids, SESSION_KEY, hops_data, onion_message=True)
 
         self.assertEqual(len(packet.to_bytes()), ONION_MESSAGE_LARGE_SIZE + 66)
 
         hops_data = hops_data_for_message('A' * ONION_MESSAGE_LARGE_SIZE)  # does not fit in ONION_MESSAGE_LARGE_SIZE
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
+
         with self.assertRaises(InvalidPayloadSize):
             new_onion_packet(blinded_node_ids, SESSION_KEY, hops_data, onion_message=True)
 
@@ -261,7 +260,7 @@ class TestOnionMessage(ElectrumTestCase):
             )
         payment_path_pubkeys = blinded_node_ids + blinded_path_blinded_ids
         hop_shared_secrets, _ = get_shared_secrets_along_route(payment_path_pubkeys, SESSION_KEY)
-        encrypt_onionmsg_tlv_hops_data(hops_data, hop_shared_secrets)
+        encrypt_hops_recipient_data('onionmsg_tlv', hops_data, hop_shared_secrets)
         packet = new_onion_packet(payment_path_pubkeys, SESSION_KEY, hops_data, onion_message=True)
         self.assertEqual(packet.to_bytes(), ONION_MESSAGE_PACKET)