soutade/iptogeo
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Tak in account server disconnect in Python test class

Browse Source
This commit is contained in:
Grégory Soutadé 2016-02-04 20:39:50 +01:00
parent ef5176e7ec
commit 1b72bc86f2
1 changed files with 15 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tests/iptogeo.py
View File

@ -1,5 +1,5 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- coding: utf-8 # -*- coding: utf-8 -*-
import socket import socket
import struct import struct
@ -32,8 +32,11 @@ class IPToGeo(object):
self._remote_port = remote_port self._remote_port = remote_port
self._timeout = timeout self._timeout = timeout
self._create_socket()
def _create_socket(self):
self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._socket.connect((remote_addr, remote_port)) self._socket.connect((self._remote_addr, self._remote_port))
def _create_request(self, int_ip): def _create_request(self, int_ip):
packet = '' packet = ''
@ -53,7 +56,7 @@ class IPToGeo(object):
if magic != IPToGeo.MAGIC: if magic != IPToGeo.MAGIC:
raise IPToGeoException('Invalid magic %08x' % (magic)) raise IPToGeoException('Invalid magic %08x' % (magic))
if err == 6: return (ipv4, None) # IP not found if err == IPToGeo.IP_NOT_FOUND: return (ipv4, None) # IP not found
if err != 0: if err != 0:
raise IPToGeoException(IPToGeo.ERRORS[err]) raise IPToGeoException(IPToGeo.ERRORS[err])
@ -69,9 +72,16 @@ class IPToGeo(object):
int_ip |= splitted_ip[3] << 0 int_ip |= splitted_ip[3] << 0
packet = self._create_request(int_ip) packet = self._create_request(int_ip)
self._socket.send(packet) try:
self._socket.send(packet)
except IOError, e:
# Give another chance (we may have been disconnected due to timeout
self._create_socket()
self._socket.send(packet)
packet = self._socket.recv(IPToGeo.PACKET_SIZE) packet = self._socket.recv(IPToGeo.PACKET_SIZE)
(ip, country_code) = self._check_request(packet) (ip, country_code) = self._check_request(packet)
if country_code:
# convert to string
country_code = '%c%c' % (country_code[0], country_code[1])
return (ip, country_code) return (ip, country_code)