Added unit tests for exercising EC2MultiRegionAddressTranslater.translate

tolbertam · tolbertam · commit 641c561cdce6 · 2015-01-12T15:51:26.000-06:00
diff --git a/driver-core/src/main/java/com/datastax/driver/core/policies/EC2MultiRegionAddressTranslater.java b/driver-core/src/main/java/com/datastax/driver/core/policies/EC2MultiRegionAddressTranslater.java
@@ -65,6 +65,11 @@ public EC2MultiRegionAddressTranslater() {
         }
     }
 
+    @VisibleForTesting
+    EC2MultiRegionAddressTranslater(DirContext ctx) {
+        this.ctx = ctx;
+    }
+
     @Override
     public InetSocketAddress translate(InetSocketAddress socketAddress) {
         InetAddress address = socketAddress.getAddress();
diff --git a/driver-core/src/test/java/com/datastax/driver/core/policies/EC2MultiRegionAddressTranslaterTest.java b/driver-core/src/test/java/com/datastax/driver/core/policies/EC2MultiRegionAddressTranslaterTest.java
@@ -1,13 +1,65 @@
 package com.datastax.driver.core.policies;
 
+import org.testng.annotations.Test;
+
+import javax.naming.NamingException;
+import javax.naming.directory.BasicAttributes;
+import javax.naming.directory.InitialDirContext;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 
-import org.testng.annotations.Test;
-
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.*;
 
 public class EC2MultiRegionAddressTranslaterTest {
+
+    @Test(groups = "unit")
+    public void should_return_same_address_when_no_entry_found() throws Exception {
+        InitialDirContext mock = mock(InitialDirContext.class);
+        when(mock.getAttributes(anyString(), any(String[].class)))
+                .thenReturn(new BasicAttributes());
+        EC2MultiRegionAddressTranslater translater = new EC2MultiRegionAddressTranslater(mock);
+
+        InetSocketAddress address = new InetSocketAddress("192.0.2.5", 9042);
+        assertThat(translater.translate(address)).isEqualTo(address);
+    }
+
+    @Test(groups = "unit")
+    public void should_return_same_address_when_exception_encountered() throws Exception {
+        InitialDirContext mock = mock(InitialDirContext.class);
+        when(mock.getAttributes(anyString(), any(String[].class)))
+                .thenThrow(new NamingException("Problem resolving address (not really)."));
+        EC2MultiRegionAddressTranslater translater = new EC2MultiRegionAddressTranslater(mock);
+
+        InetSocketAddress address = new InetSocketAddress("192.0.2.5", 9042);
+        assertThat(translater.translate(address)).isEqualTo(address);
+    }
+
+    @Test(groups = "unit")
+    public void should_return_new_address_when_match_found() throws Exception {
+        InetSocketAddress expectedAddress = new InetSocketAddress("54.32.55.66", 9042);
+
+        InitialDirContext mock = mock(InitialDirContext.class);
+        when(mock.getAttributes("5.2.0.192.in-addr.arpa", new String[] { "PTR" }))
+                .thenReturn(new BasicAttributes("PTR", expectedAddress.getHostName()));
+        EC2MultiRegionAddressTranslater translater = new EC2MultiRegionAddressTranslater(mock);
+
+        InetSocketAddress address = new InetSocketAddress("192.0.2.5", 9042);
+        assertThat(translater.translate(address)).isEqualTo(expectedAddress);
+    }
+
+    @Test(groups = "unit")
+    public void should_close_context_when_closed() throws Exception {
+        InitialDirContext mock = mock(InitialDirContext.class);
+        EC2MultiRegionAddressTranslater translater = new EC2MultiRegionAddressTranslater(mock);
+
+        // ensure close has not been called to this point.
+        verify(mock, times(0)).close();
+        translater.close();
+        // ensure close is closed.
+        verify(mock).close();
+    }
+
     @Test(groups = "unit")
     public void should_build_reversed_domain_name_for_ip_v4() throws Exception {
         InetAddress address = InetAddress.getByName("192.0.2.5");

Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,11 @@ public EC2MultiRegionAddressTranslater() {`
`65`	`65`	`}`
`66`	`66`	`}`
`67`	`67`
	`68`	`+ @VisibleForTesting`
	`69`	`+ EC2MultiRegionAddressTranslater(DirContext ctx) {`
	`70`	`+ this.ctx = ctx;`
	`71`	`+ }`
	`72`	`+`
`68`	`73`	`@Override`
`69`	`74`	`public InetSocketAddress translate(InetSocketAddress socketAddress) {`
`70`	`75`	`InetAddress address = socketAddress.getAddress();`