Skip to content

Commit f865546

Browse files
iKapitonauiKapitonau
committed
Formatting
1 parent 9616d45 commit f865546

2 files changed

Lines changed: 67 additions & 17 deletions

File tree

‎sdk/python/src/dstack_sdk/dstack_client.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ def replay_rtmrs(self) -> Dict[int, str]:
151151
rtmrs[idx] = replay_rtmr(history)
152152
return rtmrs
153153

154+
154155
class SignResponse(BaseModel):
155156
signature: str
156157
signature_chain: List[str]
@@ -165,6 +166,7 @@ def decode_signature_chain(self) -> List[bytes]:
165166
def decode_public_key(self) -> bytes:
166167
return bytes.fromhex(self.public_key)
167168

169+
168170
class VerifyResponse(BaseModel):
169171
valid: bool
170172

@@ -422,7 +424,9 @@ async def sign(self, algorithm: str, data: str | bytes) -> SignResponse:
422424
"""Signs data using a derived key."""
423425
data_bytes = data.encode() if isinstance(data, str) else data
424426
if algorithm == "secp256k1_prehashed" and len(data_bytes) != 32:
425-
raise ValueError(f"Pre-hashed signing requires a 32-byte digest, but received {len(data_bytes)} bytes")
427+
raise ValueError(
428+
f"Pre-hashed signing requires a 32-byte digest, but received {len(data_bytes)} bytes"
429+
)
426430

427431
hex_data = binascii.hexlify(data_bytes).decode()
428432
payload = {"algorithm": algorithm, "data": hex_data}

‎sdk/python/tests/test_client.py‎

Lines changed: 62 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424
def test_sync_client_get_key():
2525
client = DstackClient()
26-
result = client.get_key() # Test default algorithm (secp256k1)
26+
result = client.get_key() # Test default algorithm (secp256k1)
2727
assert isinstance(result, GetKeyResponse)
2828
assert isinstance(result.decode_key(), bytes)
2929
assert len(result.decode_key()) == 32
@@ -33,7 +33,7 @@ def test_sync_client_get_key():
3333
assert isinstance(result_ed, GetKeyResponse)
3434
assert len(result_ed.decode_key()) == 32
3535

36-
with pytest.raises(Exception): # Assuming unsupported algo raises error
36+
with pytest.raises(Exception): # Assuming unsupported algo raises error
3737
client.get_key(algorithm="rsa")
3838

3939

@@ -77,7 +77,7 @@ def check_info_response(result: InfoResponse):
7777
@pytest.mark.asyncio
7878
async def test_async_client_get_key():
7979
client = AsyncDstackClient()
80-
result = await client.get_key() # Test default algorithm (secp256k1)
80+
result = await client.get_key() # Test default algorithm (secp256k1)
8181
assert isinstance(result, GetKeyResponse)
8282
assert isinstance(result.decode_key(), bytes)
8383
assert len(result.decode_key()) == 32
@@ -87,7 +87,7 @@ async def test_async_client_get_key():
8787
assert isinstance(result_ed, GetKeyResponse)
8888
assert len(result_ed.decode_key()) == 32
8989

90-
with pytest.raises(Exception): # Assuming unsupported algo raises error
90+
with pytest.raises(Exception): # Assuming unsupported algo raises error
9191
await client.get_key(algorithm="rsa")
9292

9393

@@ -277,9 +277,11 @@ def test_emit_event_validation():
277277
client.emit_event("", "payload")
278278
assert "event name cannot be empty" in str(exc_info.value)
279279

280+
280281
SIGN_TEST_DATA = b"Test message for signing"
281282
SIGN_BAD_DATA = b"This is not the original message"
282283

284+
283285
def test_sync_sign_verify_ed25519():
284286
client = DstackClient()
285287
algo = "ed25519"
@@ -289,25 +291,41 @@ def test_sync_sign_verify_ed25519():
289291
assert len(sign_resp.decode_public_key()) > 0
290292
assert len(sign_resp.signature_chain) > 0
291293

292-
verify_resp = client.verify(algo, SIGN_TEST_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
294+
verify_resp = client.verify(
295+
algo,
296+
SIGN_TEST_DATA,
297+
sign_resp.decode_signature(),
298+
sign_resp.decode_public_key(),
299+
)
293300
assert isinstance(verify_resp, VerifyResponse)
294301
assert verify_resp.valid is True
295302

296-
verify_bad = client.verify(algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
303+
verify_bad = client.verify(
304+
algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key()
305+
)
297306
assert verify_bad.valid is False
298307

308+
299309
def test_sync_sign_verify_secp256k1():
300310
client = DstackClient()
301311
algo = "secp256k1"
302312
sign_resp = client.sign(algo, SIGN_TEST_DATA)
303313
assert isinstance(sign_resp, SignResponse)
304314

305-
verify_resp = client.verify(algo, SIGN_TEST_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
315+
verify_resp = client.verify(
316+
algo,
317+
SIGN_TEST_DATA,
318+
sign_resp.decode_signature(),
319+
sign_resp.decode_public_key(),
320+
)
306321
assert verify_resp.valid is True
307322

308-
verify_bad = client.verify(algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
323+
verify_bad = client.verify(
324+
algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key()
325+
)
309326
assert verify_bad.valid is False
310327

328+
311329
def test_sync_sign_verify_secp256k1_prehashed():
312330
client = DstackClient()
313331
algo = "secp256k1_prehashed"
@@ -317,20 +335,26 @@ def test_sync_sign_verify_secp256k1_prehashed():
317335
sign_resp = client.sign(algo, digest)
318336
assert isinstance(sign_resp, SignResponse)
319337

320-
verify_resp = client.verify(algo, digest, sign_resp.decode_signature(), sign_resp.decode_public_key())
338+
verify_resp = client.verify(
339+
algo, digest, sign_resp.decode_signature(), sign_resp.decode_public_key()
340+
)
321341
assert verify_resp.valid is True
322342

323343
bad_digest = hashlib.sha256(SIGN_BAD_DATA).digest()
324-
verify_bad = client.verify(algo, bad_digest, sign_resp.decode_signature(), sign_resp.decode_public_key())
344+
verify_bad = client.verify(
345+
algo, bad_digest, sign_resp.decode_signature(), sign_resp.decode_public_key()
346+
)
325347
assert verify_bad.valid is False
326348

349+
327350
def test_sync_sign_prehashed_length_error():
328351
client = DstackClient()
329352
algo = "secp256k1_prehashed"
330353
with pytest.raises(ValueError) as excinfo:
331354
client.sign(algo, b"too short")
332355
assert "32-byte digest" in str(excinfo.value)
333356

357+
334358
@pytest.mark.asyncio
335359
async def test_async_sign_verify_ed25519():
336360
client = AsyncDstackClient()
@@ -340,25 +364,41 @@ async def test_async_sign_verify_ed25519():
340364
assert len(sign_resp.decode_signature()) > 0
341365
assert len(sign_resp.decode_public_key()) > 0
342366

343-
verify_resp = await client.verify(algo, SIGN_TEST_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
367+
verify_resp = await client.verify(
368+
algo,
369+
SIGN_TEST_DATA,
370+
sign_resp.decode_signature(),
371+
sign_resp.decode_public_key(),
372+
)
344373
assert verify_resp.valid is True
345374

346-
verify_bad = await client.verify(algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
375+
verify_bad = await client.verify(
376+
algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key()
377+
)
347378
assert verify_bad.valid is False
348379

380+
349381
@pytest.mark.asyncio
350382
async def test_async_sign_verify_secp256k1():
351383
client = AsyncDstackClient()
352384
algo = "secp256k1"
353385
sign_resp = await client.sign(algo, SIGN_TEST_DATA)
354386
assert isinstance(sign_resp, SignResponse)
355387

356-
verify_resp = await client.verify(algo, SIGN_TEST_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
388+
verify_resp = await client.verify(
389+
algo,
390+
SIGN_TEST_DATA,
391+
sign_resp.decode_signature(),
392+
sign_resp.decode_public_key(),
393+
)
357394
assert verify_resp.valid is True
358395

359-
verify_bad = await client.verify(algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key())
396+
verify_bad = await client.verify(
397+
algo, SIGN_BAD_DATA, sign_resp.decode_signature(), sign_resp.decode_public_key()
398+
)
360399
assert verify_bad.valid is False
361400

401+
362402
@pytest.mark.asyncio
363403
async def test_async_sign_verify_secp256k1_prehashed():
364404
client = AsyncDstackClient()
@@ -368,13 +408,18 @@ async def test_async_sign_verify_secp256k1_prehashed():
368408
sign_resp = await client.sign(algo, digest)
369409
assert isinstance(sign_resp, SignResponse)
370410

371-
verify_resp = await client.verify(algo, digest, sign_resp.decode_signature(), sign_resp.decode_public_key())
411+
verify_resp = await client.verify(
412+
algo, digest, sign_resp.decode_signature(), sign_resp.decode_public_key()
413+
)
372414
assert verify_resp.valid is True
373415

374416
bad_digest = hashlib.sha256(SIGN_BAD_DATA).digest()
375-
verify_bad = await client.verify(algo, bad_digest, sign_resp.decode_signature(), sign_resp.decode_public_key())
417+
verify_bad = await client.verify(
418+
algo, bad_digest, sign_resp.decode_signature(), sign_resp.decode_public_key()
419+
)
376420
assert verify_bad.valid is False
377421

422+
378423
@pytest.mark.asyncio
379424
async def test_async_sign_prehashed_length_error():
380425
client = AsyncDstackClient()
@@ -383,6 +428,7 @@ async def test_async_sign_prehashed_length_error():
383428
await client.sign(algo, b"too short")
384429
assert "32-byte digest" in str(excinfo.value)
385430

431+
386432
# Test deprecated TappdClient
387433
def test_tappd_client_deprecated():
388434
"""Test that TappdClient shows deprecation warning."""

0 commit comments

Comments
 (0)