@@ -1453,9 +1453,11 @@ def delete_password(self, servicename, username):
1453
1453
def test_trinoquery_heartbeat_success (mock_requests , sample_post_response_data , sample_get_response_data ):
1454
1454
"""Test that heartbeat is sent periodically and does not stop on success."""
1455
1455
head_call_count = 0
1456
+
1456
1457
def fake_head (url , timeout = 10 ):
1457
1458
nonlocal head_call_count
1458
1459
head_call_count += 1
1460
+
1459
1461
class Resp :
1460
1462
status_code = 200
1461
1463
return Resp ()
@@ -1470,6 +1472,7 @@ class Resp:
1470
1472
http_scheme = "http" ,
1471
1473
)
1472
1474
query = TrinoQuery (request = req , query = "SELECT 1" , heartbeat_interval = 0.1 )
1475
+
1473
1476
def finish_query (* args , ** kwargs ):
1474
1477
query ._finished = True
1475
1478
return []
@@ -1481,6 +1484,7 @@ def finish_query(*args, **kwargs):
1481
1484
query ._stop_heartbeat ()
1482
1485
assert head_call_count >= 2
1483
1486
1487
+
1484
1488
@mock .patch ("trino.client.TrinoRequest.http" )
1485
1489
def test_trinoquery_heartbeat_failure_stops (mock_requests , sample_post_response_data , sample_get_response_data ):
1486
1490
"""Test that heartbeat stops after 3 consecutive failures."""
@@ -1506,6 +1510,7 @@ class Resp:
1506
1510
assert not query ._heartbeat_enabled
1507
1511
query ._stop_heartbeat ()
1508
1512
1513
+
1509
1514
@mock .patch ("trino.client.TrinoRequest.http" )
1510
1515
def test_trinoquery_heartbeat_404_405_stops (mock_requests , sample_post_response_data , sample_get_response_data ):
1511
1516
"""Test that heartbeat stops if server returns 404 or 405."""
@@ -1532,13 +1537,16 @@ class Resp:
1532
1537
assert not query ._heartbeat_enabled
1533
1538
query ._stop_heartbeat ()
1534
1539
1540
+
1535
1541
@mock .patch ("trino.client.TrinoRequest.http" )
1536
1542
def test_trinoquery_heartbeat_stops_on_finish (mock_requests , sample_post_response_data , sample_get_response_data ):
1537
1543
"""Test that heartbeat stops when the query is finished."""
1538
1544
head_call_count = 0
1545
+
1539
1546
def fake_head (url , timeout = 10 ):
1540
1547
nonlocal head_call_count
1541
1548
head_call_count += 1
1549
+
1542
1550
class Resp :
1543
1551
status_code = 200
1544
1552
return Resp ()
@@ -1563,13 +1571,16 @@ class Resp:
1563
1571
# Heartbeat should have stopped after query finished
1564
1572
assert head_call_count >= 1
1565
1573
1574
+
1566
1575
@mock .patch ("trino.client.TrinoRequest.http" )
1567
1576
def test_trinoquery_heartbeat_stops_on_cancel (mock_requests , sample_post_response_data , sample_get_response_data ):
1568
1577
"""Test that heartbeat stops when the query is cancelled."""
1569
1578
head_call_count = 0
1579
+
1570
1580
def fake_head (url , timeout = 10 ):
1571
1581
nonlocal head_call_count
1572
1582
head_call_count += 1
1583
+
1573
1584
class Resp :
1574
1585
status_code = 200
1575
1586
return Resp ()
0 commit comments