--- a/common/core/src/main/java/com/redhat/thermostat/common/storage/MongoStorage.java	Mon Jun 18 21:44:07 2012 +0200
+++ b/common/core/src/main/java/com/redhat/thermostat/common/storage/MongoStorage.java	Tue Jun 19 13:51:52 2012 +0200
@@ -110,7 +110,9 @@
 
     private BasicDBObject getAgentQueryKeyFromChunkOrGlobalAgent(Chunk chunk) {
         BasicDBObject queryKey = getAgentQueryKeyFromGlobalAgent();
-        if (queryKey == null && chunk.get(Key.AGENT_ID) != null) {
+        if (queryKey != null) {
+            return queryKey;
+        } else if (chunk.get(Key.AGENT_ID) != null) {
             return new BasicDBObject(KEY_AGENT_ID, chunk.get(Key.AGENT_ID));
         } else {
             return null;

--- a/common/core/src/test/java/com/redhat/thermostat/common/storage/MongoStorageTest.java	Mon Jun 18 21:44:07 2012 +0200
+++ b/common/core/src/test/java/com/redhat/thermostat/common/storage/MongoStorageTest.java	Tue Jun 19 13:51:52 2012 +0200
@@ -52,6 +52,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
+import java.util.UUID;
 
 import org.junit.After;
 import org.junit.Before;
@@ -359,4 +360,31 @@
         verify(gridFS).createFile(same(dataStream), eq("test"));
         verify(gridFSFile).save();
     }
+
+    @Test
+    public void verifyPutChunkUsesCorrectChunkAgent() throws Exception {
+        PowerMockito.whenNew(Mongo.class).withParameterTypes(MongoURI.class).withArguments(any(MongoURI.class)).thenReturn(m);
+        MongoStorage storage = makeStorage();
+        Chunk chunk1 = new Chunk(testCategory, false);
+        chunk1.put(Key.AGENT_ID, "123");
+        storage.putChunk(chunk1);
+        ArgumentCaptor<DBObject> dbobj = ArgumentCaptor.forClass(DBObject.class);
+        verify(testCollection).insert(dbobj.capture());
+        DBObject val = dbobj.getValue();
+        assertEquals("123", val.get("agent-id"));
+    }
+
+    @Test
+    public void verifyPutChunkUsesCorrectGlobalAgent() throws Exception {
+        PowerMockito.whenNew(Mongo.class).withParameterTypes(MongoURI.class).withArguments(any(MongoURI.class)).thenReturn(m);
+        MongoStorage storage = makeStorage();
+        storage.setAgentId(new UUID(1, 2));
+        Chunk chunk1 = new Chunk(testCategory, false);
+        chunk1.put(Key.AGENT_ID, "123");
+        storage.putChunk(chunk1);
+        ArgumentCaptor<DBObject> dbobj = ArgumentCaptor.forClass(DBObject.class);
+        verify(testCollection).insert(dbobj.capture());
+        DBObject val = dbobj.getValue();
+        assertEquals(new UUID(1, 2).toString(), val.get("agent-id"));
+    }
 }