Fix graphite reconnection

dropwizard · Mar 17, 2015 · dd0fca2 · dd0fca2 · analytically · Mar 18, 2015
1 parent 2f4d95f
commit dd0fca2
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 1 deletion.
diff --git a/metrics-graphite/src/main/java/com/codahale/metrics/graphite/Graphite.java b/metrics-graphite/src/main/java/com/codahale/metrics/graphite/Graphite.java
@@ -104,7 +104,7 @@ public Graphite(InetSocketAddress address, SocketFactory socketFactory, Charset
 
     @Override
     public void connect() throws IllegalStateException, IOException {
-        if (socket != null) {
+        if (isConnected()) {
             throw new IllegalStateException("Already connected");
         }
         InetSocketAddress address = this.address;

diff --git a/metrics-graphite/src/test/java/com/codahale/metrics/graphite/GraphiteTest.java b/metrics-graphite/src/test/java/com/codahale/metrics/graphite/GraphiteTest.java
@@ -12,6 +12,7 @@
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.net.UnknownHostException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Fail.failBecauseExceptionWasNotThrown;
@@ -32,6 +33,32 @@ public class GraphiteTest {
 
     @Before
     public void setUp() throws Exception {
+        final AtomicBoolean connected = new AtomicBoolean(true);
+        final AtomicBoolean closed = new AtomicBoolean(false);
+
+        when(socket.isConnected()).thenAnswer(new Answer<Boolean>() {
+            @Override
+            public Boolean answer(InvocationOnMock invocation) throws Throwable {
+                return connected.get();
+            }
+        });
+
+        when(socket.isClosed()).thenAnswer(new Answer<Boolean>() {
+            @Override
+            public Boolean answer(InvocationOnMock invocation) throws Throwable {
+                return closed.get();
+            }
+        });
+
+        doAnswer(new Answer<Void>() {
+            @Override
+            public Void answer(InvocationOnMock invocation) throws Throwable {
+                connected.set(false);
+                closed.set(true);
+                return null;
+            }
+        }).when(socket).close();
+
         when(socket.getOutputStream()).thenReturn(output);
 
         // Mock behavior of socket.getOutputStream().close() calling socket.close();

diff --git a/metrics-graphite/src/test/java/com/codahale/metrics/graphite/PickledGraphiteTest.java b/metrics-graphite/src/test/java/com/codahale/metrics/graphite/PickledGraphiteTest.java
@@ -1,6 +1,7 @@
 package com.codahale.metrics.graphite;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Fail.failBecauseExceptionWasNotThrown;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Mockito.*;
@@ -25,6 +26,7 @@
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.nio.charset.Charset;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class PickledGraphiteTest {
     private final SocketFactory socketFactory = mock(SocketFactory.class);
@@ -49,6 +51,32 @@ public class PickledGraphiteTest {
 
     @Before
     public void setUp() throws Exception {
+        final AtomicBoolean connected = new AtomicBoolean(true);
+        final AtomicBoolean closed = new AtomicBoolean(false);
+
+        when(socket.isConnected()).thenAnswer(new Answer<Boolean>() {
+            @Override
+            public Boolean answer(InvocationOnMock invocation) throws Throwable {
+                return connected.get();
+            }
+        });
+
+        when(socket.isClosed()).thenAnswer(new Answer<Boolean>() {
+            @Override
+            public Boolean answer(InvocationOnMock invocation) throws Throwable {
+                return closed.get();
+            }
+        });
+
+        doAnswer(new Answer<Void>() {
+            @Override
+            public Void answer(InvocationOnMock invocation) throws Throwable {
+                connected.set(false);
+                closed.set(true);
+                return null;
+            }
+        }).when(socket).close();
+
         when(socket.getOutputStream()).thenReturn(output);
 
         // Mock behavior of socket.getOutputStream().close() calling socket.close();
@@ -130,6 +158,18 @@ public void sanitizesValues() throws Exception {
             .isEqualTo("name value-woo 100\n");
     }
 
+    @Test
+    public void doesNotAllowDoubleConnections() throws Exception {
+        graphite.connect();
+        try {
+            graphite.connect();
+            failBecauseExceptionWasNotThrown(IllegalStateException.class);
+        } catch (IllegalStateException e) {
+            assertThat(e.getMessage())
+                    .isEqualTo("Already connected");
+        }
+    }
+
     String unpickleOutput() throws Exception {
         StringBuilder results = new StringBuilder();