Added item handling bits.

albertson/base.py

@@ -1,4 +1,7 @@
+from datetime import datetime
+
 import boto
+from boto.dynamodb.exceptions import DynamoDBKeyNotFoundError
 
 
 class CounterPool(object):
@@ -111,12 +114,66 @@ def get_table(self):
         Hook point for overriding how the CounterPool transforms table_name
         into a boto DynamoDB Table object.
         '''
-        try:
-            table = self.conn.get_table(self.get_table_name())
-        except boto.exception.DynamoDBResponseError:
-            if self.auto_create_table:
-                table = self.create_table()
-            else:
-                raise
+        if hasattr(self, '_table'):
+            table = self._table
+        else:
+            try:
+                table = self.conn.get_table(self.get_table_name())
+            except boto.exception.DynamoDBResponseError:
+                if self.auto_create_table:
+                    table = self.create_table()
+                else:
+                    raise
+
+            self._table = table
 
         return table
+
+    def create_item(self, hash_key, start=0):
+        '''
+        Hook point for overriding how the CouterPool creates a DynamoDB item
+        for a given counter when an existing item can't be found.
+        '''
+        table = self.get_table()
+        now = datetime.utcnow().replace(microsecond=0).isoformat()
+
+        item = table.new_item(
+            hash_key=hash_key,
+            attrs={
+                'created_on': now,
+                'modified_on': now,
+                'count': start,
+            }
+        )
+
+        return item
+
+    def get_item(self, hash_key, start=0):
+        '''
+        Hook point for overriding how the CouterPool fetches a DynamoDB item
+        for a given counter.
+        '''
+        table = self.get_table()
+
+        try:
+            item = table.get_item(hash_key=hash_key)
+        except DynamoDBKeyNotFoundError:
+            item = None
+
+        if item is None:
+            item = self.create_item(hash_key=hash_key, start=start)
+
+        return item
+
+
+    def get_counter(self, name):
+        '''
+        Gets the DynamoDB item behind a counter and ties it to a Counter
+        instace.
+        '''
+
+
+class Counter(object):
+    '''
+    Interface to individual counters.
+    '''

tests/base.py

@@ -1,6 +1,8 @@
+from datetime import datetime
 import unittest
 
 import boto
+from boto.dynamodb.exceptions import DynamoDBKeyNotFoundError
 
 from mock import MagicMock, sentinel
 
@@ -10,6 +12,8 @@
 
 from .dynamodb_utils import dynamo_cleanup, DynamoDeleteMixin
 
+ISO_FORMAT = '%Y-%m-%dT%H:%M:%S'
+
 
 class BaseCounterPoolTests(DynamoDeleteMixin, unittest.TestCase):
 
@@ -181,3 +185,52 @@ def test_get_missing_table_with_auto_create(self):
 
         pool.create_table.assert_called_with()
         self.assertEquals(expected, result)
+
+    @dynamo_cleanup
+    def test_create_item(self):
+        hash_key = 'test'
+        table = self.get_table()
+        pool = self.get_pool(auto_create_table=True)
+        now = datetime.utcnow().replace(microsecond=0)
+
+        expected = {
+            'counter_name': hash_key,
+            'count': 0,
+        }
+        result = pool.create_item(hash_key=hash_key)
+
+        self.assertDictContainsSubset(expected, result)
+
+        created_offset = datetime.strptime(result['created_on'], ISO_FORMAT) - now
+        modified_offset = datetime.strptime(result['modified_on'], ISO_FORMAT) - now
+
+        self.assertLess(created_offset.seconds, 2)
+        self.assertGreaterEqual(created_offset.seconds, 0)
+        self.assertLess(modified_offset.seconds, 2)
+        self.assertGreaterEqual(modified_offset.seconds, 0)
+
+        with self.assertRaises(DynamoDBKeyNotFoundError):
+            table.get_item(hash_key=hash_key, consistent_read=True)
+
+    @dynamo_cleanup
+    def test_get_missing_item(self):
+        hash_key = 'test'
+        pool = self.get_pool(auto_create_table=True)
+        pool.create_item = MagicMock(name='create_item')
+
+        expected = sentinel.item_return
+        pool.create_item.return_value = expected
+
+        result = pool.get_item(hash_key)
+
+        pool.create_item.assert_called_with(hash_key=hash_key, start=0)
+        self.assertEquals(expected, result)
+
+    def test_table_caching(self):
+        pool = self.get_pool()
+        pool._table = sentinel.cached_table
+
+        expected = sentinel.cached_table
+        result = pool.get_table()
+
+        self.assertEqual(expected, result)