-
Notifications
You must be signed in to change notification settings - Fork 75
Expand file tree
/
Copy pathtest_connection_reuse.py
More file actions
202 lines (159 loc) · 7.9 KB
/
test_connection_reuse.py
File metadata and controls
202 lines (159 loc) · 7.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
# SPDX-FileCopyrightText: © 2025 Phala Network <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import unittest.mock
import pytest
from dstack_sdk import AsyncDstackClient
from dstack_sdk import DstackClient
class TestConnectionReuse:
"""Test TCP connection reuse functionality."""
@pytest.mark.asyncio
async def test_async_context_manager_reuses_client(self):
"""Test that async context manager creates and reuses a single client."""
client = AsyncDstackClient()
# Verify client is None initially
assert client._client is None
async with client:
# Verify client is created when entering context
assert client._client is not None
first_client = client._client
# Make multiple calls and verify same client is used
with unittest.mock.patch.object(client._client, "post") as mock_post:
# Mock successful response
mock_response = unittest.mock.Mock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {"key": "test", "signature_chain": []}
mock_post.return_value = mock_response
# Make first call
await client.get_key("test1")
# Make second call
await client.get_key("test2")
# Verify same client instance was used for both calls
assert client._client is first_client
assert mock_post.call_count == 2
# Verify client is cleaned up after exiting context
assert client._client is None
def test_sync_context_manager_reuses_client(self):
"""Test that sync context manager creates and reuses a single client."""
client = DstackClient()
# Verify sync client is None initially
assert client.async_client._sync_client is None
with client:
# Verify client is created when entering context
assert client.async_client._sync_client is not None
first_client = client.async_client._sync_client
# Make multiple calls and verify same client is used
with unittest.mock.patch.object(
client.async_client._sync_client, "post"
) as mock_post:
# Mock successful response
mock_response = unittest.mock.Mock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {"key": "test", "signature_chain": []}
mock_post.return_value = mock_response
# Make first call
client.get_key("test1")
# Make second call
client.get_key("test2")
# Verify same client instance was used for both calls
assert client.async_client._sync_client is first_client
assert mock_post.call_count == 2
# Verify client is cleaned up after exiting context
assert client.async_client._sync_client is None
@pytest.mark.asyncio
async def test_async_without_context_manager_reuses_client(self):
"""Test that without context manager, clients are still reused."""
client = AsyncDstackClient()
with unittest.mock.patch("httpx.AsyncClient") as mock_async_client_class:
# Mock the context manager behavior
mock_client_instance = unittest.mock.AsyncMock()
mock_response = unittest.mock.Mock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {"key": "test", "signature_chain": []}
mock_client_instance.post = unittest.mock.AsyncMock(
return_value=mock_response
)
mock_async_client_class.return_value = mock_client_instance
# Make two calls
await client.get_key("test1")
await client.get_key("test2")
# Verify that the client was created and reused (not created twice)
assert mock_async_client_class.call_count == 1
def test_sync_without_context_manager_reuses_client(self):
"""Test that without context manager, clients are still reused."""
client = DstackClient()
with unittest.mock.patch("httpx.Client") as mock_client_class:
# Mock the context manager behavior
mock_client_instance = unittest.mock.Mock()
mock_response = unittest.mock.Mock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {"key": "test", "signature_chain": []}
mock_client_instance.post.return_value = mock_response
mock_client_class.return_value = mock_client_instance
# Make two calls
client.get_key("test1")
client.get_key("test2")
# Verify that the client was created and reused (not created twice)
assert mock_client_class.call_count == 1
@pytest.mark.asyncio
async def test_async_context_manager_with_real_requests(self):
"""Test async context manager with real requests to ensure connection reuse."""
client = AsyncDstackClient()
async with client:
# Make multiple requests - these should reuse the same connection
try:
result1 = await client.info()
result2 = await client.info()
# Both calls should succeed (assuming simulator is running)
assert result1 is not None
assert result2 is not None
assert hasattr(result1, "app_id")
assert hasattr(result2, "app_id")
except Exception:
# If simulator is not running, that's expected
# The important thing is that the context manager worked
pass
def test_sync_context_manager_with_real_requests(self):
"""Test sync context manager with real requests to ensure connection reuse."""
client = DstackClient()
with client:
# Make multiple requests - these should reuse the same connection
try:
result1 = client.info()
result2 = client.info()
# Both calls should succeed (assuming simulator is running)
assert result1 is not None
assert result2 is not None
assert hasattr(result1, "app_id")
assert hasattr(result2, "app_id")
except Exception:
# If simulator is not running, that's expected
# The important thing is that the context manager worked
pass
@pytest.mark.asyncio
async def test_async_nested_context_managers(self):
"""Test that nested async context managers work correctly."""
client = AsyncDstackClient()
async with client:
first_client = client._client
assert first_client is not None
# Nested context should use the same client
async with client:
assert client._client is first_client
# After nested exit, client should still be active
assert client._client is first_client
# After outer exit, client should be None
assert client._client is None
def test_sync_nested_context_managers(self):
"""Test that nested sync context managers work correctly."""
client = DstackClient()
with client:
first_client = client.async_client._sync_client
assert first_client is not None
# Nested context should use the same client
with client:
assert client.async_client._sync_client is first_client
# After nested exit, client should still be active
assert client.async_client._sync_client is first_client
# After outer exit, client should be None
assert client.async_client._sync_client is None