Tweak functionality and add tests

tests/test_dynamodbstreams/test_dynamodbstreams.py

@@ -1,10 +1,12 @@
 from __future__ import unicode_literals, print_function
 
+from nose.tools import assert_raises
+
 import boto3
 from moto import mock_dynamodb2, mock_dynamodbstreams
 
 
-class TestClass():
+class TestCore():
     stream_arn = None
     mocks = []
     
@@ -84,7 +86,7 @@ class TestClass():
         resp = conn.get_shard_iterator(
             StreamArn=self.stream_arn,
             ShardId=shard_id,
-            ShardIteratorType='TRIM_HORIZON'
+            ShardIteratorType='LATEST'
         )
         iterator_id = resp['ShardIterator']
 
@@ -137,3 +139,96 @@ class TestClass():
         # empty
         resp = conn.get_records(ShardIterator=resp['NextShardIterator'])
         assert len(resp['Records']) == 0
+
+
+class TestEdges():
+    mocks = []
+
+    def setup(self):
+        self.mocks = [mock_dynamodb2(), mock_dynamodbstreams()]
+        for m in self.mocks:
+            m.start()
+
+    def teardown(self):
+        for m in self.mocks:
+            m.stop()
+
+
+    def test_enable_stream_on_table(self):
+        conn = boto3.client('dynamodb', region_name='us-east-1')
+        resp = conn.create_table(
+            TableName='test-streams',
+            KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'}],
+            AttributeDefinitions=[{'AttributeName': 'id',
+                                   'AttributeType': 'S'}],
+            ProvisionedThroughput={'ReadCapacityUnits': 1,
+                                   'WriteCapacityUnits': 1}
+        )
+        assert 'StreamSpecification' not in resp['TableDescription']
+        
+        resp = conn.update_table(
+            TableName='test-streams',
+            StreamSpecification={
+                'StreamViewType': 'KEYS_ONLY'
+            }
+        )
+        assert 'StreamSpecification' in resp['TableDescription']
+        assert resp['TableDescription']['StreamSpecification'] == {
+            'StreamEnabled': True,
+            'StreamViewType': 'KEYS_ONLY'
+        }
+        assert 'LatestStreamLabel' in resp['TableDescription']
+
+        # now try to enable it again
+        with assert_raises(ValueError):
+            resp = conn.update_table(
+                TableName='test-streams',
+                StreamSpecification={
+                    'StreamViewType': 'OLD_IMAGES'
+                }
+            )
+        
+    def test_stream_with_range_key(self):
+        dyn = boto3.client('dynamodb', region_name='us-east-1')
+
+        resp = dyn.create_table(
+            TableName='test-streams',
+            KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'},
+                       {'AttributeName': 'color', 'KeyType': 'RANGE'}],
+            AttributeDefinitions=[{'AttributeName': 'id',
+                                   'AttributeType': 'S'},
+                                  {'AttributeName': 'color',
+                                   'AttributeType': 'S'}],
+            ProvisionedThroughput={'ReadCapacityUnits': 1,
+                                   'WriteCapacityUnits': 1},
+            StreamSpecification={
+                'StreamViewType': 'NEW_IMAGES'
+            }
+        )
+        stream_arn = resp['TableDescription']['LatestStreamArn']
+
+        streams = boto3.client('dynamodbstreams', region_name='us-east-1')
+        resp = streams.describe_stream(StreamArn=stream_arn)
+        shard_id = resp['StreamDescription']['Shards'][0]['ShardId']
+
+        resp = streams.get_shard_iterator(
+            StreamArn=stream_arn,
+            ShardId=shard_id,
+            ShardIteratorType='LATEST'
+        )
+        iterator_id = resp['ShardIterator']
+
+        dyn.put_item(
+            TableName='test-streams',
+            Item={'id': {'S': 'row1'}, 'color': {'S': 'blue'}}
+        )
+        dyn.put_item(
+            TableName='test-streams',
+            Item={'id': {'S': 'row2'}, 'color': {'S': 'green'}}
+        )
+
+        resp = streams.get_records(ShardIterator=iterator_id)
+        assert len(resp['Records']) == 2
+        assert resp['Records'][0]['eventName'] == 'INSERT'
+        assert resp['Records'][1]['eventName'] == 'INSERT'
+