flyte-sdk/tests/user_api/test_init_api_key.py at c1a7fb45d42852b64413d4dd9cbe577b390cf59b · flyteorg/flyte-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from unittest.mock import AsyncMock, patch

import pytest

from flyte import _initialize as init_module
from flyte._initialize import init_from_api_key
from flyte.errors import InitializationError


class TestInitFromApiKey:
    @pytest.fixture(autouse=True)
    def reset_global_state(self):
        init_module._init_config = None
        yield
        init_module._init_config = None

    @pytest.mark.asyncio
    async def test_init_from_api_key_no_key_no_env(self, monkeypatch):
        monkeypatch.delenv("FLYTE_API_KEY", raising=False)
        with pytest.raises(InitializationError, match="API key must be provided"):
            await init_from_api_key.aio()

    @patch("flyte._initialize.init")
    @patch("flyte.remote._client.auth._auth_utils.decode_api_key")
    @patch("flyte._utils.sanitize_endpoint")
    @pytest.mark.asyncio
    async def test_init_from_api_key_with_key(self, mock_sanitize, mock_decode, mock_init):
        mock_decode.return_value = ("test.endpoint.com", "client-id", "client-secret", "my-org")
        mock_sanitize.return_value = "https://test.endpoint.com"
        mock_init.aio = AsyncMock()

        await init_from_api_key.aio(api_key="encoded-key", project="proj", domain="dev")

        mock_decode.assert_called_once_with("encoded-key")
        mock_init.aio.assert_called_once()
        call_kwargs = mock_init.aio.call_args[1]
        assert call_kwargs["project"] == "proj"
        assert call_kwargs["domain"] == "dev"
        assert call_kwargs["auth_type"] == "ClientSecret"

    @patch("flyte._initialize.init")
    @patch("flyte.remote._client.auth._auth_utils.decode_api_key")
    @patch("flyte._utils.sanitize_endpoint")
    @pytest.mark.asyncio
    async def test_init_from_api_key_reads_env(self, mock_sanitize, mock_decode, mock_init, monkeypatch):
        monkeypatch.setenv("FLYTE_API_KEY", "env-key")
        mock_decode.return_value = ("endpoint.com", "cid", "csecret", "org")
        mock_sanitize.return_value = "https://endpoint.com"
        mock_init.aio = AsyncMock()

        await init_from_api_key.aio()

        mock_decode.assert_called_once_with("env-key")

    @patch("flyte._initialize.init")
    @patch("flyte.remote._client.auth._auth_utils.decode_api_key")
    @patch("flyte._utils.sanitize_endpoint")
    @pytest.mark.asyncio
    async def test_init_from_api_key_tls_defaults(self, mock_sanitize, mock_decode, mock_init):
        mock_decode.return_value = ("test.endpoint.com", "client-id", "client-secret", "my-org")
        mock_sanitize.return_value = "https://test.endpoint.com"
        mock_init.aio = AsyncMock()

        await init_from_api_key.aio(api_key="encoded-key")

        call_kwargs = mock_init.aio.call_args[1]
        assert call_kwargs["insecure"] is False
        assert call_kwargs["insecure_skip_verify"] is False
        assert call_kwargs["ca_cert_file_path"] is None

    @patch("flyte._initialize.init")
    @patch("flyte.remote._client.auth._auth_utils.decode_api_key")
    @patch("flyte._utils.sanitize_endpoint")
    @pytest.mark.asyncio
    async def test_init_from_api_key_forwards_tls_options(self, mock_sanitize, mock_decode, mock_init):
        mock_decode.return_value = ("test.endpoint.com", "client-id", "client-secret", "my-org")
        mock_sanitize.return_value = "https://test.endpoint.com"
        mock_init.aio = AsyncMock()

        await init_from_api_key.aio(api_key="encoded-key", insecure_skip_verify=True, ca_cert_file_path="/path/ca.pem")

        call_kwargs = mock_init.aio.call_args[1]
        assert call_kwargs["insecure"] is False
        assert call_kwargs["insecure_skip_verify"] is True
        assert call_kwargs["ca_cert_file_path"] == "/path/ca.pem"


class TestInitPassthrough:
    @pytest.fixture(autouse=True)
    def reset_global_state(self):
        init_module._init_config = None
        yield
        init_module._init_config = None

    @patch("flyte._initialize.init")
    @pytest.mark.asyncio
    async def test_init_passthrough_basic(self, mock_init):
        mock_init.aio = AsyncMock()
        from flyte._initialize import init_passthrough

        result = await init_passthrough.aio(endpoint="my.endpoint.com", project="proj", domain="dev")

        mock_init.aio.assert_called_once()
        call_kwargs = mock_init.aio.call_args[1]
        assert call_kwargs["auth_type"] == "Passthrough"
        assert call_kwargs["endpoint"] == "my.endpoint.com"
        assert result["endpoint"] == "my.endpoint.com"

    @patch("flyte._initialize.init")
    @pytest.mark.asyncio
    async def test_init_passthrough_reads_env_endpoint(self, mock_init, monkeypatch):
        mock_init.aio = AsyncMock()
        monkeypatch.setenv("_U_EP_OVERRIDE", "env-endpoint:8080")
        from flyte._initialize import init_passthrough

        await init_passthrough.aio()

        mock_init.aio.assert_called_once()
        call_kwargs = mock_init.aio.call_args[1]
        assert call_kwargs["endpoint"] == "env-endpoint:8080"