import pytest

from src.backend.connectors.snowflake_db import constants, utils

MOCK_PEM_KEY = """-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASC...
-----END PRIVATE KEY-----"""
MOCK_BASE64_KEY = "TUlJRXZRSUJBREFORmdrcWhraU5n..."  # Shortened for test


def test_get_private_key_content_empty():
    with pytest.raises(ValueError):
        utils.get_private_key_content("")


def test_get_private_key_content_pem():
    result = utils.get_private_key_content(MOCK_PEM_KEY)
    assert result == MOCK_PEM_KEY.rstrip("\n ").encode("utf-8").rstrip(b"\n")


def test_get_private_key_content_file(tmp_path):
    key_path = tmp_path / "key.pem"
    key_content = MOCK_PEM_KEY + "\n"  # Add trailing newline
    key_path.write_text(key_content)
    result = utils.get_private_key_content(str(key_path))
    assert result == key_content.rstrip("\n ").encode("utf-8").rstrip(b"\n")


def test_get_private_key_content_base64():
    result = utils.get_private_key_content(MOCK_BASE64_KEY)
    assert result.startswith(constants.PRIVATE_KEY_HEADER)
    assert result.endswith(constants.PRIVATE_KEY_FOOTER.rstrip(b"\n"))