Test mongo conditional upsert

'''
Test for mongo conditional upsert
http://stackoverflow.com/a/19214951/60422

Created on 2013-10-06

@author: rakslice
'''
import pymongo
import unittest

# Obviously this doesn't test for race conditions.

db = pymongo.MongoClient().mongo_conditional_update
collection = db.test

assert isinstance(collection, pymongo.collection.Collection)


def update_if_stale(key, new_value, new_time):
    collection.update({'key': key,
                       'update_time': {'$lt': new_time}
                       },
                      {'$set': {'value': new_value,
                                'update_time': new_time
                                }
                       }
                      )


def insert_if_missing(key, new_value, new_time):
    collection.update({'key': key},
                      {'$setOnInsert': {'value': new_value,
                                        'update_time': new_time
                                        }
                       },
                      upsert=True
                      )


def update_key(key, new_value, new_time):
    insert_if_missing(key, new_value, new_time)
    update_if_stale(key, new_value, new_time)


def get_key(key):
    result = collection.find_one({'key': key})
    if result is None:
        return result
    return result["value"]


def key_count(key):
    result = list(collection.find({'key': key}))
    return len(result)


class TestConditionalUpdate(unittest.TestCase):
    def setUp(self):
        collection.drop()

    def testPreviouslyDidNotExist(self):
        update_key("a", "foo", 10)
        self.assertEqual("foo", get_key("a"))

    def testPreviouslyHadStaleValue(self):
        collection.insert({"key": "a", "value": "joe", "update_time": 5})

        assert "joe" == get_key("a")

        update_key("a", "foo", 10)

        assert "foo" == get_key("a")
        assert 1 == key_count("a")

    def testPreviouslyHadNewValue(self):
        collection.insert({"key": "a", "value": "joe", "update_time": 15})

        assert "joe" == get_key("a")

        update_key("a", "foo", 10)

        assert "joe" == get_key("a")
        assert 1 == key_count("a")