add partition key to env var

layercake/tests/test_dynamodb.py

@@ -6,10 +6,11 @@ from botocore.exceptions import ClientError
 
 from layercake.dateutils import ttl
 from layercake.dynamodb import (
+    ComposeKey,
     DynamoDBCollection,
     DynamoDBPersistenceLayer,
-    Key,
     KeyPair,
+    PartitionKey,
     TransactItems,
     serialize,
 )
@@ -31,12 +32,22 @@ def test_serialize():
     }
 
 
-def test_key():
-    assert Key(('122', 'abc'), prefix='schedules') == 'schedules#122#abc'
+def test_composekey():
+    assert ComposeKey(('122', 'abc'), prefix='schedules') == 'schedules#122#abc'
+    assert ComposeKey(('122', 'abc')) == '122#abc'
+    assert ComposeKey('122') == '122'
+
+
+def test_partitionkey():
+    assert PartitionKey('123') == {'id': '123'}
+    assert PartitionKey('123').expr_attr_name() == {'#pk': 'id'}
+    assert PartitionKey('123').expr_attr_values() == {':pk': '123'}
 
 
 def test_keypair():
     assert KeyPair('123', 'abc') == {'id': '123', 'sk': 'abc'}
+    assert KeyPair('123', 'abc').expr_attr_name() == {'#pk': 'id', '#sk': 'sk'}
+    assert KeyPair('123', 'abc').expr_attr_values() == {':pk': '123', ':sk': 'abc'}
 
 
 def test_transact_write_items(
@@ -65,7 +76,7 @@ def test_collection_get_item(
 ):
     collect = DynamoDBCollection(dynamodb_persistence_layer)
     data_notfound = collect.get_item(
-        key=KeyPair(
+        KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
             sk='tenant',
         ),
@@ -76,9 +87,9 @@ def test_collection_get_item(
 
     # This item was added from seeds
     data = collect.get_item(
-        key=KeyPair(
+        KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
-            sk=Key('sergio@somosbeta.com.br', prefix='emails'),
+            sk=ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
         ),
         default={},
     )
@@ -102,18 +113,18 @@ def test_collection_put_item(
     collect = DynamoDBCollection(dynamodb_persistence_layer)
 
     assert collect.put_item(
-        key=KeyPair(
+        KeyPair(
             '5OxmMjL-ujoR5IMGegQz',
-            Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
+            ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
         ),
         name='Beta Educação',
         ttl=ttl(days=3),
     )
 
     data = collect.get_item(
-        key=KeyPair(
+        KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
-            sk=Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
+            sk=ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
         ),
     )
 
@@ -131,8 +142,23 @@ def test_collection_delete_item(
 
     # This item was added from seeds
     assert collect.delete_item(
-        key=KeyPair(
+        KeyPair(
             '5OxmMjL-ujoR5IMGegQz',
-            Key('sergio@somsbeta.com.br', prefix='emails'),
+            ComposeKey('sergio@somsbeta.com.br', prefix='emails'),
         )
     )
+
+
+def test_collection_get_items(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+
+    # This item was added from seeds
+    data = collect.get_items(
+        PartitionKey(
+            ComposeKey('5OxmMjL-ujoR5IMGegQz', prefix='logs'),
+        ),
+    )
+    assert len(data['items']) == 2