flyte-sdk/tests/flyte/remote/test_token_client.py at ebea97a8c256cd6faac4b54d9908c1e6837348ad · 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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
"""Tests for how the OAuth token/device-code endpoints handle a body that is not JSON.

Both endpoints are specified to answer in JSON, but the SDK talks to them through whatever a
deployment puts in the way: a load balancer's HTML 502 page, a proxy's plain-text "Internal
Server Error", an SSO interstitial. Reading such a body used to raise `json.JSONDecodeError`
straight out of the error branch that was about to raise a perfectly good `AuthenticationError`
(FLYTE-SDK-60).
"""

from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest

from flyte.remote._client.auth._token_client import (
    GrantType,
    _body_snippet,
    _json_object_or_none,
    get_device_code,
    get_token,
)
from flyte.remote._client.auth.errors import AuthenticationError, AuthenticationPending

# What a proxy in front of a broken IDP actually returns -- the FLYTE-SDK-60 event.
NGINX_502 = "<html><head><title>502 Bad Gateway</title></head><body>502 Bad Gateway</body></html>"


def _session(response: httpx.Response) -> MagicMock:
    session = MagicMock()
    session.post = AsyncMock(return_value=response)
    return session


class TestJsonObjectOrNone:
    @pytest.mark.parametrize(
        "response, expected",
        [
            (httpx.Response(200, json={"access_token": "t"}), {"access_token": "t"}),
            (httpx.Response(500, text=NGINX_502), None),
            (httpx.Response(500, text="Internal Server Error"), None),
            (httpx.Response(200, text=""), None),
            # Valid JSON that is not an object: `"error" in j` would answer by substring on a
            # bare string, which is not the membership test the caller means.
            (httpx.Response(400, json="error"), None),
            (httpx.Response(400, json=["error"]), None),
        ],
    )
    def test_only_json_objects_survive(self, response, expected):
        assert _json_object_or_none(response) == expected


class TestBodySnippet:
    def test_reports_content_type_and_body(self):
        snippet = _body_snippet(httpx.Response(500, text="boom", headers={"content-type": "text/plain"}))
        assert "text/plain" in snippet
        assert "boom" in snippet

    def test_empty_body_is_said_to_be_empty(self):
        assert "empty body" in _body_snippet(httpx.Response(500, text=""))

    def test_long_body_is_truncated(self):
        snippet = _body_snippet(httpx.Response(500, text="x" * 5000), limit=50)
        assert len(snippet) < 200
        assert snippet.endswith("...'")


class TestGetTokenNonJsonBody:
    @pytest.mark.asyncio
    async def test_error_status_with_html_body_reports_the_status(self):
        """The failure branch already had the right error to raise; it just had to get there."""
        with pytest.raises(AuthenticationError) as excinfo:
            await get_token("https://idp.example.com/token", _session(httpx.Response(500, text=NGINX_502)))

        assert "Status Code (500)" in str(excinfo.value)

    @pytest.mark.asyncio
    async def test_error_status_with_empty_body_reports_the_status(self):
        with pytest.raises(AuthenticationError) as excinfo:
            await get_token("https://idp.example.com/token", _session(httpx.Response(503, text="")))

        assert "Status Code (503)" in str(excinfo.value)

    @pytest.mark.asyncio
    async def test_success_status_with_non_json_body_names_the_endpoint(self):
        """A 200 carrying a login page is not an SDK bug, and must not read as one."""
        with pytest.raises(AuthenticationError) as excinfo:
            await get_token("https://idp.example.com/token", _session(httpx.Response(200, text=NGINX_502)))

        message = str(excinfo.value)
        assert "https://idp.example.com/token" in message
        assert "not an access token" in message

    @pytest.mark.asyncio
    async def test_success_status_without_access_token_is_reported(self):
        with pytest.raises(AuthenticationError) as excinfo:
            await get_token("https://idp.example.com/token", _session(httpx.Response(200, json={"scope": "all"})))

        assert "not an access token" in str(excinfo.value)


class TestGetTokenStillWorks:
    @pytest.mark.asyncio
    async def test_access_token_is_returned(self):
        response = httpx.Response(200, json={"access_token": "abc", "refresh_token": "r", "expires_in": 3600})

        access, refresh, expires = await get_token("https://idp.example.com/token", _session(response))

        assert (access, refresh, expires) == ("abc", "r", 3600)

    @pytest.mark.asyncio
    async def test_missing_refresh_token_is_fine(self):
        response = httpx.Response(200, json={"access_token": "abc", "expires_in": 3600})

        access, refresh, expires = await get_token("https://idp.example.com/token", _session(response))

        assert (access, refresh, expires) == ("abc", None, 3600)

    @pytest.mark.asyncio
    @pytest.mark.parametrize("err", ["authorization_pending", "slow_down"])
    async def test_device_flow_pending_still_raises_authentication_pending(self, err):
        """The JSON error branch is the one that keeps the device-code poll loop alive."""
        response = httpx.Response(400, json={"error": err})

        with pytest.raises(AuthenticationPending):
            await get_token(
                "https://idp.example.com/token",
                _session(response),
                grant_type=GrantType.DEVICE_CODE,
                device_code="dc",
            )


class TestGetDeviceCodeNonJsonBody:
    @pytest.mark.asyncio
    async def test_error_status_does_not_crash_building_its_own_message(self):
        """`Reason {resp.json()}` was interpolated into the error it was raising."""
        with pytest.raises(AuthenticationError) as excinfo:
            await get_device_code(
                "https://idp.example.com/device", "client", _session(httpx.Response(502, text=NGINX_502))
            )

        assert "Status Code 502" in str(excinfo.value)

    @pytest.mark.asyncio
    async def test_success_status_with_non_json_body_names_the_endpoint(self):
        with pytest.raises(AuthenticationError) as excinfo:
            await get_device_code(
                "https://idp.example.com/device", "client", _session(httpx.Response(200, text=NGINX_502))
            )

        assert "https://idp.example.com/device" in str(excinfo.value)

    @pytest.mark.asyncio
    async def test_valid_device_code_response_is_parsed(self):
        response = httpx.Response(
            200,
            json={
                "device_code": "dc",
                "user_code": "UC",
                "verification_uri": "https://idp.example.com/activate",
                "expires_in": 600,
                "interval": 5,
            },
        )

        result = await get_device_code("https://idp.example.com/device", "client", _session(response))

        assert result.device_code == "dc"
        assert result.interval == 5