diff --git a/src/documents.js b/src/documents.js
index e79506f..0e9d2fd 100644
--- a/src/documents.js
+++ b/src/documents.js
@@ -10,9 +10,11 @@ const DocumentStore = async ({ OpLog, Database, ipfs, identity, databaseId, acce
    * @returns {string} The hash of the new oplog entry.
    */
   const put = async (doc) => {
-    if (!doc[indexBy]) { throw new Error(`The provided document doesn't contain field '${indexBy}'`) }
+    const key = doc[indexBy]
     
-    return addOperation({ op: 'PUT', key: doc[indexBy], value: doc })
+    if (!key) { throw new Error(`The provided document doesn't contain field '${indexBy}'`) }
+    
+    return addOperation({ op: 'PUT', key, value: doc })
   }
 
   /**
@@ -34,29 +36,25 @@ const DocumentStore = async ({ OpLog, Database, ipfs, identity, databaseId, acce
    * @returns {Object} The doc corresponding to key or null.
    */
   const get = async (key) => {
-    for await (const entry of iterator()) {
-      const { key: k, value } = entry
-      
-      if (key === k) {
-        return value  
+    for await (const doc of iterator()) {
+      if (key === doc[indexBy]) {
+        return doc  
       }
     }
-    
-    return null
   }
   
   /**
    * Queries the document store for documents matching mapper filters.
    *
-   * @param {function(Object)} mapper A function for querying for specific results.
+   * @param {function(Object)} findFn A function for querying for specific results.
    * @returns {Array} Found documents.
    */
-  const query = async (mapper) => {
+  const query = async (findFn) => {
     const results = []
     
-    for await (const entry of iterator()) {
-      if (Object.values(entry).find(mapper)) {
-        results.push(entry.value)  
+    for await (const doc of iterator()) {
+      if (findFn(doc)) {
+        results.push(doc)
       }
     }
 
@@ -69,7 +67,7 @@ const DocumentStore = async ({ OpLog, Database, ipfs, identity, databaseId, acce
       const { op, key, value } = entry.payload
       if (op === 'PUT' && !keys[key]) {
         keys[key] = true
-        yield { key, value }
+        yield value
       } else if (op === 'DEL' && !keys[key]) {
         keys[key] = true
       }
diff --git a/test/documents.spec.js b/test/documents.spec.js
index 1f783bf..9cf74b7 100644
--- a/test/documents.spec.js
+++ b/test/documents.spec.js
@@ -105,35 +105,37 @@ Object.keys(testAPIs).forEach((IPFS) => {
     describe('using database', () => {
       it('gets a document', async () => {
          const key = 'hello world 1'
+         
+         const expected = { _id: key, msg: 'writing 1 to db1' }
 
-         await db1.put({ _id: key, doc: 'writing 1 to db1' })
+         await db1.put(expected)
          
          const doc = await db1.get(key)
-         strictEqual(doc._id, key)
+         deepStrictEqual(doc, expected)
       })
       
       it('deletes a document', async () => {
          const key = 'hello world 1'
          
-         await db1.put({ _id: key, doc: 'writing 1 to db1' })
+         await db1.put({ _id: key, msg: 'writing 1 to db1' })
          await db1.del(key)
          
          const doc = await db1.get(key)
-         strictEqual(doc, null)
+         strictEqual(doc, undefined)
       })
 
       it('queries a document', async () => {
-          const expected = { _id: 'hello world 1', doc: 'writing new 1 to db1', views: 10 }
+          const expected = { _id: 'hello world 1', msg: 'writing new 1 to db1', views: 10 }
           
-          const doc3 = { _id: 'hello world 3', doc: 'writing 3 to db1', views: 12 }
-          
-          await db1.put({ _id: 'hello world 1', doc: 'writing 1 to db1', views: 10 })
-          await db1.put({ _id: 'hello world 2', doc: 'writing 2 to db1', views: 5 })
-          await db1.put({ _id: 'hello world 3', doc: 'writing 3 to db1', views: 12 })
+          await db1.put({ _id: 'hello world 1', msg: 'writing 1 to db1', views: 10 })
+          await db1.put({ _id: 'hello world 2', msg: 'writing 2 to db1', views: 5 })
+          await db1.put({ _id: 'hello world 3', msg: 'writing 3 to db1', views: 12 })
           await db1.del('hello world 3')          
           await db1.put(expected)
           
-          deepStrictEqual(await db1.query((e) => e.views > 5), [expected])
+          const findFn = (doc) => doc.views > 5
+          
+          deepStrictEqual(await db1.query(findFn), [expected])
       })
     })
   })