From 0435548ae881d69cd7c549460781d459a4782371 Mon Sep 17 00:00:00 2001
From: Omer Katz <omer.katz@kcg.tech>
Date: Thu, 3 Oct 2024 11:55:25 +0300
Subject: [PATCH] Added unit tests.

---
 t/unit/test_connection.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/t/unit/test_connection.py b/t/unit/test_connection.py
index 25cee80b1..35abb1c77 100644
--- a/t/unit/test_connection.py
+++ b/t/unit/test_connection.py
@@ -752,6 +752,41 @@ def test_connection_failover_with_total_timeout(self):
                 conn.default_channel
             assert conn._establish_connection.call_count == 2
 
+    def test_connection_timeout_with_errback(self):
+        errback = Mock()
+        with Connection(
+            ['server1', 'server2'],
+            transport=TimeoutingTransport,
+            connect_timeout=1,
+            transport_options={
+                'connect_retries_timeout': 2,
+                'interval_start': 0,
+                'interval_step': 0,
+                'errback': errback
+            },
+        ) as conn:
+            with pytest.raises(OperationalError):
+                conn.default_channel
+
+        errback.assert_called()
+
+    def test_connection_timeout_with_callback(self):
+        callback = Mock()
+        with Connection(
+            ['server1', 'server2'],
+            transport=TimeoutingTransport,
+            connect_timeout=1,
+            transport_options={
+                'connect_retries_timeout': 2,
+                'interval_start': 0,
+                'interval_step': 0,
+                'callback': callback
+            },
+        ) as conn:
+            with pytest.raises(OperationalError):
+                conn.default_channel
+
+        callback.assert_called()
 
 class test_Connection_with_transport_options: