Skip to content

Commit 1e9feeb

Browse files
quettabitquettabit
committed
add tests
1 parent 0b7215c commit 1e9feeb

1 file changed

Lines changed: 196 additions & 0 deletions

File tree

‎tests/test_stream_ops.py‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import asyncio
2+
import base64
3+
import os
24
import time
35
import uuid
46
from datetime import timedelta
@@ -12,6 +14,7 @@
1214
Batching,
1315
CommandRecord,
1416
Compression,
17+
EncryptionAlgorithm,
1518
Endpoints,
1619
FencingTokenMismatchError,
1720
ReadLimit,
@@ -844,3 +847,196 @@ async def test_compression_roundtrip_session(
844847
await basin.delete_stream(stream_name)
845848
finally:
846849
await s2.delete_basin(basin_name)
850+
851+
852+
def _make_key(key_type: str) -> bytes | str:
853+
raw = os.urandom(32)
854+
return raw if key_type == "bytes" else base64.b64encode(raw).decode()
855+
856+
857+
@pytest.mark.stream
858+
@pytest.mark.parametrize(
859+
"cipher",
860+
[EncryptionAlgorithm.AEGIS_256, EncryptionAlgorithm.AES_256_GCM],
861+
)
862+
class TestEncryption:
863+
@pytest.mark.parametrize("key_type", ["bytes", "str"])
864+
async def test_encryption_roundtrip_unary(
865+
self,
866+
access_token: str,
867+
endpoints: Endpoints | None,
868+
cipher: EncryptionAlgorithm,
869+
key_type: str,
870+
):
871+
async with S2(access_token, endpoints=endpoints) as s2:
872+
basin_name = f"test-py-sdk-{uuid.uuid4().hex[:8]}"
873+
await s2.create_basin(
874+
name=basin_name,
875+
config=BasinConfig(
876+
stream_cipher=cipher,
877+
default_stream_config=StreamConfig(
878+
storage_class=StorageClass.STANDARD
879+
),
880+
),
881+
)
882+
try:
883+
basin = s2.basin(basin_name)
884+
stream_name = f"stream-{uuid.uuid4().hex[:8]}"
885+
info = await basin.create_stream(
886+
name=stream_name,
887+
config=StreamConfig(timestamping=Timestamping(uncapped=True)),
888+
)
889+
assert info.cipher == cipher
890+
891+
key = _make_key(key_type)
892+
try:
893+
stream = basin.stream(stream_name, encryption_key=key)
894+
ack = await stream.append(
895+
AppendInput(
896+
records=[
897+
Record(body=b"hello"),
898+
Record(body=b"world"),
899+
]
900+
)
901+
)
902+
assert ack.start.seq_num == 0
903+
assert ack.end.seq_num == 2
904+
905+
batch = await stream.read(start=SeqNum(0))
906+
assert len(batch.records) == 2
907+
assert batch.records[0].body == b"hello"
908+
assert batch.records[1].body == b"world"
909+
finally:
910+
await basin.delete_stream(stream_name)
911+
finally:
912+
await s2.delete_basin(basin_name)
913+
914+
@pytest.mark.parametrize("key_type", ["bytes", "str"])
915+
async def test_encryption_roundtrip_session(
916+
self,
917+
access_token: str,
918+
endpoints: Endpoints | None,
919+
cipher: EncryptionAlgorithm,
920+
key_type: str,
921+
):
922+
async with S2(access_token, endpoints=endpoints) as s2:
923+
basin_name = f"test-py-sdk-{uuid.uuid4().hex[:8]}"
924+
await s2.create_basin(
925+
name=basin_name,
926+
config=BasinConfig(
927+
stream_cipher=cipher,
928+
default_stream_config=StreamConfig(
929+
timestamping=Timestamping(mode=TimestampingMode.ARRIVAL)
930+
),
931+
),
932+
)
933+
try:
934+
basin = s2.basin(basin_name)
935+
stream_name = f"stream-{uuid.uuid4().hex[:8]}"
936+
info = await basin.create_stream(
937+
name=stream_name,
938+
config=StreamConfig(storage_class=StorageClass.STANDARD),
939+
)
940+
assert info.cipher == cipher
941+
942+
key = _make_key(key_type)
943+
try:
944+
stream = basin.stream(stream_name, encryption_key=key)
945+
async with stream.append_session() as session:
946+
ticket = await session.submit(
947+
AppendInput(records=[Record(body=b"s2" * 10240)])
948+
)
949+
ack = await ticket
950+
951+
assert ack.start.seq_num == 0
952+
assert ack.end.seq_num == 1
953+
954+
batch = await stream.read(start=SeqNum(0))
955+
assert len(batch.records) == 1
956+
assert batch.records[0].body == b"s2" * 10240
957+
finally:
958+
await basin.delete_stream(stream_name)
959+
finally:
960+
await s2.delete_basin(basin_name)
961+
962+
async def test_encrypted_stream_requires_key(
963+
self,
964+
access_token: str,
965+
endpoints: Endpoints | None,
966+
cipher: EncryptionAlgorithm,
967+
):
968+
async with S2(access_token, endpoints=endpoints) as s2:
969+
basin_name = f"test-py-sdk-{uuid.uuid4().hex[:8]}"
970+
await s2.create_basin(
971+
name=basin_name,
972+
config=BasinConfig(
973+
stream_cipher=cipher,
974+
default_stream_config=StreamConfig(
975+
storage_class=StorageClass.STANDARD
976+
),
977+
),
978+
)
979+
try:
980+
basin = s2.basin(basin_name)
981+
stream_name = f"stream-{uuid.uuid4().hex[:8]}"
982+
await basin.create_stream(
983+
name=stream_name,
984+
config=StreamConfig(timestamping=Timestamping(uncapped=True)),
985+
)
986+
try:
987+
stream = basin.stream(stream_name)
988+
with pytest.raises(S2ServerError) as append_exc:
989+
await stream.append(
990+
AppendInput(records=[Record(body=b"hello")])
991+
)
992+
assert append_exc.value.code == "invalid"
993+
assert cipher.value in str(append_exc.value)
994+
with pytest.raises(S2ServerError) as read_exc:
995+
await stream.read(start=SeqNum(0))
996+
assert read_exc.value.code == "invalid"
997+
assert cipher.value in str(read_exc.value)
998+
finally:
999+
await basin.delete_stream(stream_name)
1000+
finally:
1001+
await s2.delete_basin(basin_name)
1002+
1003+
async def test_encrypted_stream_wrong_key_read_fails(
1004+
self,
1005+
access_token: str,
1006+
endpoints: Endpoints | None,
1007+
cipher: EncryptionAlgorithm,
1008+
):
1009+
async with S2(access_token, endpoints=endpoints) as s2:
1010+
basin_name = f"test-py-sdk-{uuid.uuid4().hex[:8]}"
1011+
await s2.create_basin(
1012+
name=basin_name,
1013+
config=BasinConfig(
1014+
stream_cipher=cipher,
1015+
default_stream_config=StreamConfig(
1016+
storage_class=StorageClass.STANDARD
1017+
),
1018+
),
1019+
)
1020+
try:
1021+
basin = s2.basin(basin_name)
1022+
stream_name = f"stream-{uuid.uuid4().hex[:8]}"
1023+
await basin.create_stream(
1024+
name=stream_name,
1025+
config=StreamConfig(timestamping=Timestamping(uncapped=True)),
1026+
)
1027+
try:
1028+
key_a = os.urandom(32)
1029+
key_b = os.urandom(32)
1030+
1031+
writer = basin.stream(stream_name, encryption_key=key_a)
1032+
reader = basin.stream(stream_name, encryption_key=key_b)
1033+
1034+
await writer.append(AppendInput(records=[Record(body=b"secret")]))
1035+
1036+
with pytest.raises(S2ServerError) as exc_info:
1037+
await reader.read(start=SeqNum(0))
1038+
assert exc_info.value.code == "decryption_failed"
1039+
finally:
1040+
await basin.delete_stream(stream_name)
1041+
finally:
1042+
await s2.delete_basin(basin_name)

0 commit comments

Comments
 (0)