1414# limitations under the License.
1515import asyncio
1616import sys
17- import time
18- from asyncio .exceptions import TimeoutError , InvalidStateError
17+ from asyncio import TimeoutError , InvalidStateError
1918
2019import pytest
2120
@@ -73,10 +72,12 @@ async def test_cancellation_on_slow_response(event_loop, proxy, invalid_states_e
7372 c = AioClient (partition_aware = False )
7473 async with c .connect ("127.0.0.1" , 10802 ):
7574 cache = await c .get_cache ("test" )
76- proxy .slow_response = True
75+ proxy .discard_response = True # Simulate slow response by discarding it
76+
7777 with pytest .raises (TimeoutError ):
7878 await asyncio .wait_for (cache .put (1 , 2 ), 0.1 )
7979
80+ proxy .discard_response = False
8081 assert len (invalid_states_errors ) == 0
8182
8283
@@ -96,15 +97,14 @@ async def test_cancellation_on_disconnect(event_loop, proxy, invalid_states_erro
9697
9798class ProxyServer :
9899 """
99- Proxy for simulating slow or discarding response ignite server
100- Set `slow_response`, ` discard_response` to `True` to simulate these conditions .
100+ Proxy for simulating discarding response from ignite server
101+ Set `discard_response` to `True` to simulate this condition .
101102 Call `disconnect_peers()` in order to simulate lost connection to Ignite server.
102103 """
103104 def __init__ (self , local_host , remote_host ):
104105 self .local_host = local_host
105106 self .remote_host = remote_host
106107 self .peers = {}
107- self .slow_response = False
108108 self .discard_response = False
109109 self .server = None
110110
@@ -206,7 +206,4 @@ def data_received(self, data):
206206 if self .proxy .discard_response :
207207 return
208208
209- if self .proxy .slow_response :
210- time .sleep (0.5 )
211-
212209 self .proxy_protocol .transport .write (data )
0 commit comments