Make IPv6 opt-in an explicit choice.

3 files changed, 47 insertions, 13 deletions

diff --git a/CHANGES b/CHANGES
index 5408ee0..6882961 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,6 +1,12 @@
 Statsd Changelog
 ================
 
+Version 3.2
+-----------
+
+- Add an explicit IPv6 flag.
+
+
 Version 3.1
 -----------
 
diff --git a/statsd/client.py b/statsd/client.py
index a423dae..73e7712 100644
--- a/statsd/client.py
+++ b/statsd/client.py
@@ -131,12 +131,12 @@ class StatsClientBase(object):
 class StatsClient(StatsClientBase):
     """A client for statsd."""
 
-    def __init__(self, host='localhost', port=8125, prefix=None,
+    def __init__(self, host='localhost', port=8125, ipv6=False, prefix=None,
                  maxudpsize=512):
         """Create a new client."""
+        fam = socket.AF_INET6 if ipv6 else socket.AF_INET
         family, _, _, _, addr = socket.getaddrinfo(
-            host, port, 0, socket.SOCK_DGRAM
-        )[0]
+            host, port, fam, socket.SOCK_DGRAM)[0]
         self._addr = addr
         self._sock = socket.socket(family, socket.SOCK_DGRAM)
         self._prefix = prefix
@@ -157,10 +157,12 @@ class StatsClient(StatsClientBase):
 class TCPStatsClient(StatsClientBase):
     """TCP version of StatsClient."""
 
-    def __init__(self, host='localhost', port=8125, prefix=None, timeout=None):
+    def __init__(self, host='localhost', port=8125, ipv6=False, prefix=None,
+                 timeout=None):
         """Create a new client."""
         self._host = host
         self._port = port
+        self._ipv6 = ipv6
         self._timeout = timeout
         self._prefix = prefix
         self._sock = None
@@ -180,8 +182,9 @@ class TCPStatsClient(StatsClientBase):
         self._sock = None
 
     def connect(self):
+        fam = socket.AF_INET6 if self._ipv6 else socket.AF_INET
         family, _, _, _, addr = socket.getaddrinfo(
-            self._host, self._port, 0, socket.SOCK_STREAM)[0]
+            self._host, self._port, fam, socket.SOCK_STREAM)[0]
         self._sock = socket.socket(family, socket.SOCK_STREAM)
         self._sock.settimeout(self._timeout)
         self._sock.connect(addr)
diff --git a/statsd/tests.py b/statsd/tests.py
index a6a525c..16cf7ac 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -27,22 +27,23 @@ make_val = {
 }
 
 
-def _udp_client(prefix=None, addr=None, port=None):
+def _udp_client(prefix=None, addr=None, port=None, ipv6=False):
     if not addr:
         addr = ADDR[0]
     if not port:
         port = ADDR[1]
-    sc = StatsClient(host=addr, port=port, prefix=prefix)
+    sc = StatsClient(host=addr, port=port, prefix=prefix, ipv6=ipv6)
     sc._sock = mock.Mock()
     return sc
 
 
-def _tcp_client(prefix=None, addr=None, port=None, timeout=None):
+def _tcp_client(prefix=None, addr=None, port=None, timeout=None, ipv6=False):
     if not addr:
         addr = ADDR[0]
     if not port:
         port = ADDR[1]
-    sc = TCPStatsClient(host=addr, port=port, prefix=prefix, timeout=timeout)
+    sc = TCPStatsClient(host=addr, port=port, prefix=prefix, timeout=timeout,
+                        ipv6=ipv6)
     sc._sock = mock.Mock()
     return sc
 
@@ -210,21 +211,45 @@ def _test_ipv6(cl, proto, addr):
     _sock_check(cl._sock, 1, proto, 'foo:30|g', addr=addr)
 
 
-@mock.patch.object(random, 'random', lambda: -1)
 def test_ipv6_udp():
     """StatsClient can use to IPv6 address."""
     addr = ('::1', 8125, 0, 0)
-    cl = _udp_client(addr=addr[0])
+    cl = _udp_client(addr=addr[0], ipv6=True)
     _test_ipv6(cl, 'udp', addr)
 
-@mock.patch.object(random, 'random', lambda: -1)
+
 def test_ipv6_tcp():
     """TCPStatsClient can use to IPv6 address."""
     addr = ('::1', 8125, 0, 0)
-    cl = _tcp_client(addr=addr[0])
+    cl = _tcp_client(addr=addr[0], ipv6=True)
     _test_ipv6(cl, 'tcp', addr)
 
 
+def _test_resolution(cl, proto, addr):
+    cl.incr('foo')
+    _sock_check(cl._sock, 1, proto, 'foo:1|c', addr=addr)
+
+
+def test_ipv6_resolution_udp():
+    cl = _udp_client(addr='localhost', ipv6=True)
+    _test_resolution(cl, 'udp', ('::1', 8125, 0, 0))
+
+
+def test_ipv6_resolution_tcp():
+    cl = _tcp_client(addr='localhost', ipv6=True)
+    _test_resolution(cl, 'tcp', ('::1', 8125, 0, 0))
+
+
+def test_ipv4_resolution_udp():
+    cl = _udp_client(addr='localhost')
+    _test_resolution(cl, 'udp', ('127.0.0.1', 8125))
+
+
+def test_ipv4_resolution_tcp():
+    cl = _tcp_client(addr='localhost')
+    _test_resolution(cl, 'tcp', ('127.0.0.1', 8125))
+
+
 def _test_gauge_delta(cl, proto):
     tests = (
         (12, '+12'),