From f5f38ff23e89e924f2828c7675f427e9c985a291 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benjamindhimes@gmail.com>
Date: Thu, 22 May 2025 16:41:36 +0200
Subject: [PATCH] Add proper mock support for SubstrateInterface and
 AsyncSubstrateInterface

---
 async_substrate_interface/async_substrate.py | 22 +++++++++------
 async_substrate_interface/sync_substrate.py  | 11 ++++++--
 tests/unit_tests/test_mock.py                | 29 ++++++++++++++++++++
 3 files changed, 51 insertions(+), 11 deletions(-)
 create mode 100644 tests/unit_tests/test_mock.py

diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py
index 94bda13..312fb30 100644
--- a/async_substrate_interface/async_substrate.py
+++ b/async_substrate_interface/async_substrate.py
@@ -9,6 +9,7 @@
 import logging
 import ssl
 import time
+from unittest.mock import AsyncMock
 from hashlib import blake2b
 from typing import (
     Optional,
@@ -697,13 +698,16 @@ def __init__(
         self.chain_endpoint = url
         self.url = url
         self._chain = chain_name
-        self.ws = Websocket(
-            url,
-            options={
-                "max_size": self.ws_max_size,
-                "write_limit": 2**16,
-            },
-        )
+        if not _mock:
+            self.ws = Websocket(
+                url,
+                options={
+                    "max_size": self.ws_max_size,
+                    "write_limit": 2**16,
+                },
+            )
+        else:
+            self.ws = AsyncMock(spec=Websocket)
         self._lock = asyncio.Lock()
         self.config = {
             "use_remote_preset": use_remote_preset,
@@ -726,9 +730,11 @@ def __init__(
         self._initializing = False
         self.registry_type_map = {}
         self.type_id_to_name = {}
+        self._mock = _mock
 
     async def __aenter__(self):
-        await self.initialize()
+        if not self._mock:
+            await self.initialize()
         return self
 
     async def initialize(self):
diff --git a/async_substrate_interface/sync_substrate.py b/async_substrate_interface/sync_substrate.py
index 4c91fd2..dc8d178 100644
--- a/async_substrate_interface/sync_substrate.py
+++ b/async_substrate_interface/sync_substrate.py
@@ -3,6 +3,7 @@
 import socket
 from hashlib import blake2b
 from typing import Optional, Union, Callable, Any
+from unittest.mock import MagicMock
 
 from bt_decode import MetadataV15, PortableRegistry, decode as decode_by_type_string
 from scalecodec import (
@@ -13,7 +14,7 @@
     MultiAccountId,
 )
 from scalecodec.base import RuntimeConfigurationObject, ScaleBytes, ScaleType
-from websockets.sync.client import connect
+from websockets.sync.client import connect, ClientConnection
 from websockets.exceptions import ConnectionClosed
 
 from async_substrate_interface.const import SS58_FORMAT
@@ -522,14 +523,18 @@ def __init__(
         )
         self.metadata_version_hex = "0x0f000000"  # v15
         self.reload_type_registry()
-        self.ws = self.connect(init=True)
         self.registry_type_map = {}
         self.type_id_to_name = {}
+        self._mock = _mock
         if not _mock:
+            self.ws = self.connect(init=True)
             self.initialize()
+        else:
+            self.ws = MagicMock(spec=ClientConnection)
 
     def __enter__(self):
-        self.initialize()
+        if not self._mock:
+            self.initialize()
         return self
 
     def __del__(self):
diff --git a/tests/unit_tests/test_mock.py b/tests/unit_tests/test_mock.py
new file mode 100644
index 0000000..81a1f50
--- /dev/null
+++ b/tests/unit_tests/test_mock.py
@@ -0,0 +1,29 @@
+from websockets.exceptions import InvalidURI
+import pytest
+
+from async_substrate_interface import AsyncSubstrateInterface, SubstrateInterface
+
+
+@pytest.mark.asyncio
+async def test_async_mock():
+    ssi = AsyncSubstrateInterface("notreal")
+    assert isinstance(ssi, AsyncSubstrateInterface)
+    with pytest.raises(InvalidURI):
+        await ssi.initialize()
+    async with AsyncSubstrateInterface("notreal", _mock=True) as ssi:
+        assert isinstance(ssi, AsyncSubstrateInterface)
+    ssi = AsyncSubstrateInterface("notreal", _mock=True)
+    async with ssi:
+        pass
+
+
+def test_sync_mock():
+    with pytest.raises(InvalidURI):
+        SubstrateInterface("notreal")
+    ssi = SubstrateInterface("notreal", _mock=True)
+    assert isinstance(ssi, SubstrateInterface)
+    with pytest.raises(InvalidURI):
+        with SubstrateInterface("notreal") as ssi:
+            pass
+    with SubstrateInterface("notreal", _mock=True) as ssi:
+        assert isinstance(ssi, SubstrateInterface)