--- a/agent/src/main/java/com/redhat/thermostat/agent/storage/Category.java	Wed Feb 22 20:59:37 2012 +0100
+++ b/agent/src/main/java/com/redhat/thermostat/agent/storage/Category.java	Wed Feb 22 22:57:35 2012 +0100
@@ -37,15 +37,32 @@
 package com.redhat.thermostat.agent.storage;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Set;
 
 public class Category {
     private final String name;
     private final List<Key> keys;
     private boolean locked = false;
 
+    private ConnectionKey connectionKey;
+
+    private static Set<String> categoryNames = new HashSet<String>();
+
+    /**
+     * Creates a new Category instance with the specified name.
+     *
+     * @param name the name of the category
+     *
+     * @throws IllegalArgumentException if a Category is created with a name that has been used before
+     */
     public Category(String name) {
+        if (categoryNames.contains(name)) {
+            throw new IllegalStateException();
+        }
+        categoryNames.add(name);
         this.name = name;
         keys = new ArrayList<Key>();
     }
@@ -69,4 +86,16 @@
     public synchronized Iterator<Key> getEntryIterator() {
         return keys.iterator();
     }
+
+    public void setConnectionKey(ConnectionKey connKey) {
+        connectionKey = connKey;
+    }
+
+    public ConnectionKey getConnectionKey() {
+        return connectionKey;
+    }
+
+    public boolean hasBeenRegistered() {
+        return getConnectionKey() != null;
+    }
 }

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/agent/src/main/java/com/redhat/thermostat/agent/storage/ConnectionKey.java	Wed Feb 22 22:57:35 2012 +0100
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+package com.redhat.thermostat.agent.storage;
+
+public interface ConnectionKey {
+
+}

--- a/agent/src/main/java/com/redhat/thermostat/agent/storage/MongoStorage.java	Wed Feb 22 20:59:37 2012 +0100
+++ b/agent/src/main/java/com/redhat/thermostat/agent/storage/MongoStorage.java	Wed Feb 22 22:57:35 2012 +0100
@@ -119,7 +119,7 @@
     }
 
     @Override
-    protected void putChunkImpl(Chunk chunk) {
+    public void putChunk(Chunk chunk) {
         Category cat = chunk.getCategory();
         DBCollection coll = getCachedCollection(cat.getName());
         BasicDBObject toInsert = getAgentDBObject();
@@ -179,7 +179,7 @@
     }
 
     @Override
-    protected void updateChunkImpl(Chunk chunk) {
+    public void updateChunk(Chunk chunk) {
         Category cat = chunk.getCategory();
         DBCollection coll = getCachedCollection(cat.getName());
         BasicDBObject toUpdate = new BasicDBObject();
@@ -285,4 +285,11 @@
             coll.remove(deleteKey);
         }
     }
+
+    @Override
+    public ConnectionKey createConnectionKey(Category category) {
+        // TODO: We want to return an instance of an inner class here that carries the actual connection
+        // and replace the collectionCache. For now this is good enough though.
+        return new ConnectionKey(){};
+    }
 }

--- a/agent/src/main/java/com/redhat/thermostat/agent/storage/Storage.java	Wed Feb 22 20:59:37 2012 +0100
+++ b/agent/src/main/java/com/redhat/thermostat/agent/storage/Storage.java	Wed Feb 22 22:57:35 2012 +0100
@@ -37,18 +37,9 @@
 package com.redhat.thermostat.agent.storage;
 
 import java.net.UnknownHostException;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.UUID;
 
-import com.redhat.thermostat.backend.Backend;
-
 public abstract class Storage {
-    private Map<String, Backend> categoryMap;
-
-    public Storage() {
-        categoryMap = new HashMap<String, Backend>();
-    }
 
     public abstract void connect(String uri) throws UnknownHostException;
 
@@ -63,33 +54,19 @@
      */
     public abstract String getBackendConfig(String backendName, String configurationKey);
 
-    public final void registerCategory(Category category, Backend backend) {
-        if (categoryMap.containsKey(category.getName())) {
+    public final void registerCategory(Category category) {
+        if (category.hasBeenRegistered()) {
             throw new IllegalStateException("Category may only be associated with one backend.");
         }
-        categoryMap.put(category.getName(), backend);
-    }
-
-    public final void putChunk(Chunk chunk, Backend backend) {
-        validateChunkOrigin(chunk, backend);
-        putChunkImpl(chunk);
+        ConnectionKey connKey = createConnectionKey(category);
+        category.setConnectionKey(connKey);
     }
 
-    public final void updateChunk(Chunk chunk, Backend backend) {
-        validateChunkOrigin(chunk, backend);
-        updateChunkImpl(chunk);
-    }
+    public abstract ConnectionKey createConnectionKey(Category category);
 
-    private void validateChunkOrigin(Chunk chunk, Backend origin) {
-        Category category = chunk.getCategory();
-        if (origin != categoryMap.get(category.getName())) { // This had better be not just equivalent, but actually the same object.
-            throw new IllegalArgumentException("Invalid category-backend combination while inserting data.  Category: " + category.getName() + "  Backend: " + origin.getName());
-        }
-    }
-    
-    protected abstract void putChunkImpl(Chunk chunk);
+    public abstract void putChunk(Chunk chunk);
 
-    protected abstract void updateChunkImpl(Chunk chunk);
+    public abstract void updateChunk(Chunk chunk);
 
     /* Drop all data related to the currently running agent.
      */

--- a/agent/src/main/java/com/redhat/thermostat/backend/Backend.java	Wed Feb 22 20:59:37 2012 +0100
+++ b/agent/src/main/java/com/redhat/thermostat/backend/Backend.java	Wed Feb 22 22:57:35 2012 +0100
@@ -80,7 +80,7 @@
     public final void setStorage(Storage storage) {
         this.storage = storage;
         for (Iterator<Category> iter = getCategoryIterator(); iter.hasNext();) {
-            storage.registerCategory(iter.next(), this);
+            storage.registerCategory(iter.next());
         }
     }
 
@@ -189,10 +189,10 @@
     }
 
     public final void store(Chunk chunk) {
-        storage.putChunk(chunk, this);
+        storage.putChunk(chunk);
     }
 
     public final void update(Chunk chunk) {
-        storage.updateChunk(chunk, this);
+        storage.updateChunk(chunk);
     }
 }

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/agent/src/test/java/com/redhat/thermostat/agent/storage/CategoryTest.java	Wed Feb 22 22:57:35 2012 +0100
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+package com.redhat.thermostat.agent.storage;
+
+import static org.junit.Assert.*;
+
+import org.junit.Test;
+
+public class CategoryTest {
+
+    @Test
+    public void testHasBeenRegistered() {
+        Category category1 = new Category("testHasBeenRegistered");
+        assertFalse(category1.hasBeenRegistered());
+        ConnectionKey connKey = new ConnectionKey(){};
+        category1.setConnectionKey(connKey);
+        assertTrue(category1.hasBeenRegistered());
+    }
+
+}

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/agent/src/test/java/com/redhat/thermostat/agent/storage/MongoStorageTest.java	Wed Feb 22 22:57:35 2012 +0100
@@ -0,0 +1,53 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+package com.redhat.thermostat.agent.storage;
+
+import static org.junit.Assert.assertNotNull;
+
+import org.junit.Test;
+
+public class MongoStorageTest {
+
+    @Test
+    public void testCreateConnectionKey() {
+        MongoStorage mongoStorage = new MongoStorage();
+        Category category = new Category("test");
+        ConnectionKey connKey = mongoStorage.createConnectionKey(category);
+        assertNotNull(connKey);
+    }
+
+}

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/agent/src/test/java/com/redhat/thermostat/agent/storage/StorageTest.java	Wed Feb 22 22:57:35 2012 +0100
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2012 Red Hat, Inc.
+ *
+ * This file is part of Thermostat.
+ *
+ * Thermostat is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation; either version 2, or (at your
+ * option) any later version.
+ *
+ * Thermostat is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Thermostat; see the file COPYING.  If not see
+ * <http://www.gnu.org/licenses/>.
+ *
+ * Linking this code with other modules is making a combined work
+ * based on this code.  Thus, the terms and conditions of the GNU
+ * General Public License cover the whole combination.
+ *
+ * As a special exception, the copyright holders of this code give
+ * you permission to link this code with independent modules to
+ * produce an executable, regardless of the license terms of these
+ * independent modules, and to copy and distribute the resulting
+ * executable under terms of your choice, provided that you also
+ * meet, for each linked independent module, the terms and conditions
+ * of the license of that module.  An independent module is a module
+ * which is not derived from or based on this code.  If you modify
+ * this code, you may extend this exception to your version of the
+ * library, but you are not obligated to do so.  If you do not wish
+ * to do so, delete this exception statement from your version.
+ */
+
+package com.redhat.thermostat.agent.storage;
+
+import static org.junit.Assert.assertSame;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class StorageTest {
+
+    private Storage storage;
+    private ConnectionKey connKey;
+
+    @Before
+    public void setUp() {
+        storage = mock(Storage.class);
+        connKey = new ConnectionKey(){};
+        when(storage.createConnectionKey(any(Category.class))).thenReturn(connKey);
+    }
+
+    @After
+    public void tearDown() {
+        storage = null;
+        connKey = null;
+    }
+
+    @Test
+    public void testRegisterCategory() {
+        Category category = new Category("testRegisterCategory");
+        storage.registerCategory(category);
+
+        verify(storage).createConnectionKey(category);
+        assertSame(connKey, category.getConnectionKey());
+    }
+
+    @Test(expected=IllegalStateException.class)
+    public void testRegisterCategoryTwice() {
+
+        Category category = new Category("test");
+        storage.registerCategory(category);
+        storage.registerCategory(category);
+    }
+
+    @Test(expected=IllegalStateException.class)
+    public void testRegisterCategorySameName() {
+
+        Category category1 = new Category("test");
+        storage.registerCategory(category1);
+        Category category2 = new Category("test");
+        storage.registerCategory(category2);
+    }
+}