From 5d4f811a65062f11c262df51a6ca78214ec981ea Mon Sep 17 00:00:00 2001
From: uael <uael@google.com>
Date: Tue, 2 May 2023 05:31:52 +0000
Subject: [PATCH] smp: add simple `Session` proxy

This allow modifying the SMP behavior at runtime for testing purpose.
---
 bumble/device.py | 10 +++++++++-
 bumble/smp.py    | 10 ++++++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/bumble/device.py b/bumble/device.py
index 7448bd8f..72fd7550 100644
--- a/bumble/device.py
+++ b/bumble/device.py
@@ -23,7 +23,7 @@ import asyncio
 import logging
 from contextlib import asynccontextmanager, AsyncExitStack
 from dataclasses import dataclass
-from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Type, Union
 
 from .colors import color
 from .att import ATT_CID, ATT_DEFAULT_MTU, ATT_PDU
@@ -2206,6 +2206,14 @@ class Device(CompositeEventEmitter):
     ) -> None:
         self.smp_manager.pairing_config_factory = pairing_config_factory
 
+    @property
+    def smp_session_proxy(self) -> Type[smp.Session]:
+        return self.smp_manager.session_proxy
+
+    @smp_session_proxy.setter
+    def smp_session_proxy(self, session_proxy: Type[smp.Session]) -> None:
+        self.smp_manager.session_proxy = session_proxy
+
     async def pair(self, connection):
         return await self.smp_manager.pair(connection)
 
diff --git a/bumble/smp.py b/bumble/smp.py
index ade51ec0..f3fbf279 100644
--- a/bumble/smp.py
+++ b/bumble/smp.py
@@ -1711,6 +1711,7 @@ class Manager(EventEmitter):
     device: Device
     sessions: Dict[int, Session]
     pairing_config_factory: Callable[[Connection], PairingConfig]
+    session_proxy: Type[Session]
 
     def __init__(
         self,
@@ -1722,6 +1723,7 @@ class Manager(EventEmitter):
         self.sessions = {}
         self._ecc_key = None
         self.pairing_config_factory = pairing_config_factory
+        self.session_proxy = Session
 
     def send_command(self, connection: Connection, command: SMP_Command) -> None:
         logger.debug(
@@ -1737,7 +1739,9 @@ class Manager(EventEmitter):
             if connection.role == BT_CENTRAL_ROLE:
                 logger.warning('Remote starts pairing as Peripheral!')
             pairing_config = self.pairing_config_factory(connection)
-            session = Session(self, connection, pairing_config, is_initiator=False)
+            session = self.session_proxy(
+                self, connection, pairing_config, is_initiator=False
+            )
             self.sessions[connection.handle] = session
 
         # Parse the L2CAP payload into an SMP Command object
@@ -1762,7 +1766,9 @@ class Manager(EventEmitter):
         if connection.role != BT_CENTRAL_ROLE:
             logger.warning('Start pairing as Peripheral!')
         pairing_config = self.pairing_config_factory(connection)
-        session = Session(self, connection, pairing_config, is_initiator=True)
+        session = self.session_proxy(
+            self, connection, pairing_config, is_initiator=True
+        )
         self.sessions[connection.handle] = session
         return await session.pair()