"""Tests for prodigy_teams.key_pair — RSA key pair generation."""

from __future__ import annotations

import base64

from prodigy_teams.key_pair import generate_key_pair


class TestKeyPair:
    def test_generate_no_password(self):
        kp = generate_key_pair()
        assert kp.public_pem.startswith("-----BEGIN PUBLIC KEY-----")
        assert kp.private_pem.startswith("-----BEGIN PRIVATE KEY-----")

    def test_generate_with_password(self):
        kp = generate_key_pair(password="secret123")
        assert kp.public_pem.startswith("-----BEGIN PUBLIC KEY-----")
        assert kp.private_pem.startswith("-----BEGIN ENCRYPTED PRIVATE KEY-----")

    def test_b64_encoding(self):
        kp = generate_key_pair()
        decoded_pub = base64.b64decode(kp.public_pem_b64).decode("utf8")
        assert decoded_pub == kp.public_pem
        decoded_priv = base64.b64decode(kp.private_pem_b64).decode("utf8")
        assert decoded_priv == kp.private_pem

    def test_keys_are_different(self):
        kp1 = generate_key_pair()
        kp2 = generate_key_pair()
        assert kp1.private_pem != kp2.private_pem
        assert kp1.public_pem != kp2.public_pem
