# HG changeset patch
# User Severin Gehwolf <sgehwolf@redhat.com>
# Date 1395396906 -3600
# Node ID 0bb2a75521f3d288a007ab1c917d12773b18e31a
# Parent  0cbf519e37f0725f904d4e0432625683c2b02869
Fix MongoStorage's ConnectionListener.
Reviewed-by: neugens, omajid
Review-thread: http://icedtea.classpath.org/pipermail/thermostat/2014-March/009496.html
PR1702

diff -r 0cbf519e37f0 -r 0bb2a75521f3 storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java
--- a/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Mon Feb 24 18:34:23 2014 -0700
+++ b/storage/mongo/src/main/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorage.java	Fri Mar 21 11:15:06 2014 +0100
@@ -288,25 +288,42 @@
         this.conn = null;
     }
     
-    public MongoStorage(String url, StorageCredentials creds, SSLConfiguration sslConf) {
-        conn = new MongoConnection(url, creds, sslConf);
+    MongoStorage(MongoConnection connection) {
+        this.conn = connection;
         connectedLatch = new CountDownLatch(1);
+        
+        // We register a connection listener in order for the mongo-java-driver
+        // DB object to be valid once it's first used (that's usually in
+        // registerCategory())
         conn.addListener(new ConnectionListener() {
             @Override
             public void changed(ConnectionStatus newStatus) {
                 switch (newStatus) {
-                case DISCONNECTED:
-                    db = null;
                 case CONNECTED:
+                    // Main success entry point
                     db = conn.getDB();
                     // This is important. See comment in registerCategory().
                     connectedLatch.countDown();
-                default:
-                    // ignore other status events
+                    break;
+                case FAILED_TO_CONNECT:
+                    // Main connection-failure entry-point
+                    connectedLatch.countDown();
+                    break;
+                case CONNECTING:
+                    // no-op
+                    break;
+                case DISCONNECTED:
+                    // mark the db object invalid
+                    db = null;
+                    break;
                 }
             }
         });
     }
+    
+    public MongoStorage(String url, StorageCredentials creds, SSLConfiguration sslConf) {
+    	this(new MongoConnection(url, creds, sslConf));
+    }
 
     public <T extends Pojo> Cursor<T> executeGetCount(Category<T> category, MongoQuery<T> queryToAggregate) {
         try {
diff -r 0cbf519e37f0 -r 0bb2a75521f3 storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java
--- a/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Mon Feb 24 18:34:23 2014 -0700
+++ b/storage/mongo/src/test/java/com/redhat/thermostat/storage/mongodb/internal/MongoStorageTest.java	Fri Mar 21 11:15:06 2014 +0100
@@ -87,6 +87,8 @@
 import com.redhat.thermostat.storage.core.BackingStorage;
 import com.redhat.thermostat.storage.core.Category;
 import com.redhat.thermostat.storage.core.CategoryAdapter;
+import com.redhat.thermostat.storage.core.Connection.ConnectionListener;
+import com.redhat.thermostat.storage.core.Connection.ConnectionStatus;
 import com.redhat.thermostat.storage.core.Cursor;
 import com.redhat.thermostat.storage.core.Entity;
 import com.redhat.thermostat.storage.core.Key;
@@ -99,6 +101,7 @@
 import com.redhat.thermostat.storage.model.AggregateCount;
 import com.redhat.thermostat.storage.model.BasePojo;
 import com.redhat.thermostat.storage.model.HostInfo;
+import com.redhat.thermostat.storage.model.Pojo;
 import com.redhat.thermostat.storage.query.Expression;
 import com.redhat.thermostat.storage.query.ExpressionFactory;
 
@@ -262,6 +265,49 @@
     }
     
     @Test
+    public void verifyConnectEventSetsDb() {
+        MongoConnection mockConnection = mock(MongoConnection.class);
+        
+        // This adds a listener which we capture
+        new MongoStorage(mockConnection);
+        
+        ArgumentCaptor<ConnectionListener> captor = ArgumentCaptor.forClass(ConnectionListener.class);
+        verify(mockConnection).addListener(captor.capture());
+        ConnectionListener listener = captor.getValue();
+        assertNotNull(listener);
+        listener.changed(ConnectionStatus.CONNECTED);
+        verify(mockConnection).getDB();
+    }
+    
+    @Test(expected=NullPointerException.class)
+    public void verifyDisconnectedEventInvalidatesDb() {
+        MongoConnection mockConnection = mock(MongoConnection.class);
+        
+        // This adds a listener which we capture
+        MongoStorage storage = new MongoStorage(mockConnection);
+        
+        ArgumentCaptor<ConnectionListener> captor = ArgumentCaptor.forClass(ConnectionListener.class);
+        verify(mockConnection).addListener(captor.capture());
+        ConnectionListener listener = captor.getValue();
+        assertNotNull(listener);
+        
+        // fire a connecting event for good measure. It should do nothing.
+        listener.changed(ConnectionStatus.CONNECTING);
+        
+        listener.changed(ConnectionStatus.CONNECTED);
+        verify(mockConnection).getDB();
+        
+        // This should set the db instance to null
+        listener.changed(ConnectionStatus.DISCONNECTED);
+        
+        @SuppressWarnings("unchecked")
+        Category<FakeDataClass> cat = mock(Category.class);
+        when(cat.getDataClass()).thenReturn(FakeDataClass.class);
+        // this throws NPE due to the null DB instance
+        storage.registerCategory(cat);
+    }
+    
+    @Test
     public void testRegisterCategory() throws Exception {
         DB db = PowerMockito.mock(DB.class);
         CountDownLatch latch = new CountDownLatch(1);
@@ -618,5 +664,7 @@
         dbField.setAccessible(true);
         dbField.set(storage, db);
     }
+    
+    private static class FakeDataClass implements Pojo {};
 }