Skip to content

Commit 59ca6a6

Browse files
canniumcannium
committed
Improvements based on code review.
1 parent 91eb6d4 commit 59ca6a6

File tree

2 files changed

+49
-61
lines changed

2 files changed

+49
-61
lines changed

‎influxdb/client.py

Lines changed: 14 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
Python client for InfluxDB
44
"""
55
from collections import OrderedDict
6+
from functools import wraps
67
import json
78
import socket
89
import random
@@ -481,12 +482,11 @@ def __init__(self,
481482
use_udp=False,
482483
udp_port=4444,
483484
shuffle=True,
484-
client_base_class=InfluxDBClient,
485+
client_base_class=InfluxDBClient, # For simpler test code
485486
):
486487
self.clients = []
487488
self.bad_clients = [] # Corresponding server has failures in history
488489
self.shuffle = shuffle # if true, queries will hit servers evenly
489-
self.client_base_class = client_base_class # For simpler test code
490490
for h in hosts:
491491
self.clients.append(client_base_class(host=h[0], port=h[1],
492492
username=username,
@@ -500,13 +500,14 @@ def __init__(self,
500500
for method in dir(client_base_class):
501501
if method.startswith('_'):
502502
continue
503-
if not callable(getattr(client_base_class, method)):
503+
orig_func = getattr(client_base_class, method)
504+
if not callable(orig_func):
504505
continue
505-
setattr(self, method, self._make_func(method))
506+
setattr(self, method, self._make_func(orig_func))
506507

507-
def _make_func(self, func_name):
508-
orig_func = getattr(self.client_base_class, func_name)
508+
def _make_func(self, orig_func):
509509

510+
@wraps(orig_func)
510511
def func(*args, **kwargs):
511512
if self.shuffle:
512513
random.shuffle(self.clients)
@@ -521,21 +522,14 @@ def func(*args, **kwargs):
521522
except Exception as e:
522523
# Errors that might caused by server failure, try another
523524
bad_client = True
525+
if c in self.clients:
526+
self.clients.remove(c)
527+
self.bad_clients.append(c)
524528
finally:
525-
if bad_client:
526-
if c not in self.bad_clients:
527-
self.bad_clients.append(c)
528-
for idx, val in enumerate(self.clients):
529-
if val == c:
530-
del self.clients[idx]
531-
break
532-
else:
533-
if c not in self.clients:
534-
self.clients.append(c)
535-
for idx, val in enumerate(self.bad_clients):
536-
if val == c:
537-
del self.bad_clients[idx]
538-
break
529+
if not bad_client and c in self.bad_clients:
530+
self.bad_clients.remove(c)
531+
self.clients.append(c)
532+
539533
raise InfluxDBServerError("InfluxDB: no viable server!")
540534

541535
return func

‎tests/influxdb/client_test.py

Lines changed: 35 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -97,41 +97,41 @@ def setUp(self):
9797

9898
def test_scheme(self):
9999
cli = InfluxDBClient('host', 8086, 'username', 'password', 'database')
100-
assert cli._baseurl == 'http://host:8086'
100+
self.assertEqual('http://host:8086', cli._baseurl)
101101

102102
cli = InfluxDBClient(
103103
'host', 8086, 'username', 'password', 'database', ssl=True
104104
)
105-
assert cli._baseurl == 'https://host:8086'
105+
self.assertEqual('https://host:8086', cli._baseurl)
106106

107107
def test_dsn(self):
108108
cli = InfluxDBClient.from_DSN('influxdb://usr:pwd@host:1886/db')
109-
assert cli._baseurl == 'http://host:1886'
110-
assert cli._username == 'usr'
111-
assert cli._password == 'pwd'
112-
assert cli._database == 'db'
113-
assert cli.use_udp is False
109+
self.assertEqual('http://host:1886', cli._baseurl)
110+
self.assertEqual('usr', cli._username)
111+
self.assertEqual('pwd', cli._password)
112+
self.assertEqual('db', cli._database)
113+
self.assertFalse(cli.use_udp)
114114

115115
cli = InfluxDBClient.from_DSN('udp+influxdb://usr:pwd@host:1886/db')
116-
assert cli.use_udp is True
116+
self.assertTrue(cli.use_udp)
117117

118118
cli = InfluxDBClient.from_DSN('https+influxdb://usr:pwd@host:1886/db')
119-
assert cli._baseurl == 'https://host:1886'
119+
self.assertEqual('https://host:1886', cli._baseurl)
120120

121121
cli = InfluxDBClient.from_DSN('https+influxdb://usr:pwd@host:1886/db',
122122
**{'ssl': False})
123-
assert cli._baseurl == 'http://host:1886'
123+
self.assertEqual('http://host:1886', cli._baseurl)
124124

125125
def test_switch_database(self):
126126
cli = InfluxDBClient('host', 8086, 'username', 'password', 'database')
127127
cli.switch_database('another_database')
128-
assert cli._database == 'another_database'
128+
self.assertEqual('another_database', cli._database)
129129

130130
def test_switch_user(self):
131131
cli = InfluxDBClient('host', 8086, 'username', 'password', 'database')
132132
cli.switch_user('another_username', 'another_password')
133-
assert cli._username == 'another_username'
134-
assert cli._password == 'another_password'
133+
self.assertEqual('another_username', cli._username)
134+
self.assertEqual('another_password', cli._password)
135135

136136
def test_write(self):
137137
with requests_mock.Mocker() as m:
@@ -208,10 +208,8 @@ def test_write_points_toplevel_attributes(self):
208208
def test_write_points_batch(self):
209209
cli = InfluxDBClient('host', 8086, 'username', 'password', 'db')
210210
with _mocked_session(cli, 'post', 200, self.dummy_points):
211-
assert cli.write_points(
212-
data=self.dummy_points,
213-
batch_size=2
214-
) is True
211+
self.assertTrue(cli.write_points(data=self.dummy_points,
212+
batch_size=2))
215213

216214
def test_write_points_udp(self):
217215
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -569,21 +567,21 @@ def test_init(self):
569567
database='database',
570568
shuffle=False,
571569
client_base_class=FakeClient)
572-
assert len(cluster.clients) == 3
573-
assert len(cluster.bad_clients) == 0
570+
self.assertEqual(3, len(cluster.clients))
571+
self.assertEqual(0, len(cluster.bad_clients))
574572
for idx, client in enumerate(cluster.clients):
575-
assert client._host == self.hosts[idx][0]
576-
assert client._port == self.hosts[idx][1]
573+
self.assertEqual(self.hosts[idx][0], client._host)
574+
self.assertEqual(self.hosts[idx][1], client._port)
577575

578576
def test_one_server_fails(self):
579577
cluster = InfluxDBClusterClient(hosts=self.hosts,
580578
database='database',
581579
shuffle=False,
582580
client_base_class=FakeClient)
583581
cluster.clients[0].fail = True
584-
assert cluster.query('') == 'Success'
585-
assert len(cluster.clients) == 2
586-
assert len(cluster.bad_clients) == 1
582+
self.assertEqual('Success', cluster.query(''))
583+
self.assertEqual(2, len(cluster.clients))
584+
self.assertEqual(1, len(cluster.bad_clients))
587585

588586
def test_two_servers_fail(self):
589587
cluster = InfluxDBClusterClient(hosts=self.hosts,
@@ -592,40 +590,36 @@ def test_two_servers_fail(self):
592590
client_base_class=FakeClient)
593591
cluster.clients[0].fail = True
594592
cluster.clients[1].fail = True
595-
assert cluster.query('') == 'Success'
596-
assert len(cluster.clients) == 1
597-
assert len(cluster.bad_clients) == 2
593+
self.assertEqual('Success', cluster.query(''))
594+
self.assertEqual(1, len(cluster.clients))
595+
self.assertEqual(2, len(cluster.bad_clients))
598596

599597
def test_all_fail(self):
600598
cluster = InfluxDBClusterClient(hosts=self.hosts,
601599
database='database',
602600
shuffle=True,
603601
client_base_class=FakeClient)
604-
try:
602+
with self.assertRaises(InfluxDBServerError):
605603
cluster.query('Fail')
606-
except InfluxDBServerError:
607-
pass
608-
assert len(cluster.clients) == 0
609-
assert len(cluster.bad_clients) == 3
604+
self.assertEqual(0, len(cluster.clients))
605+
self.assertEqual(3, len(cluster.bad_clients))
610606

611607
def test_all_good(self):
612608
cluster = InfluxDBClusterClient(hosts=self.hosts,
613609
database='database',
614610
shuffle=True,
615611
client_base_class=FakeClient)
616-
assert cluster.query('') == 'Success'
617-
assert len(cluster.clients) == 3
618-
assert len(cluster.bad_clients) == 0
612+
self.assertEqual('Success', cluster.query(''))
613+
self.assertEqual(3, len(cluster.clients))
614+
self.assertEqual(0, len(cluster.bad_clients))
619615

620616
def test_recovery(self):
621617
cluster = InfluxDBClusterClient(hosts=self.hosts,
622618
database='database',
623619
shuffle=True,
624620
client_base_class=FakeClient)
625-
try:
621+
with self.assertRaises(InfluxDBServerError):
626622
cluster.query('Fail')
627-
except InfluxDBServerError:
628-
pass
629-
assert cluster.query('') == 'Success'
630-
assert len(cluster.clients) == 1
631-
assert len(cluster.bad_clients) == 2
623+
self.assertEqual('Success', cluster.query(''))
624+
self.assertEqual(1, len(cluster.clients))
625+
self.assertEqual(2, len(cluster.bad_clients))

0 commit comments

Comments
 (0)