Coverage for tests / unit / test_session.py: 100%
138 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-02 13:07 +0100
1import base64
2from unittest.mock import Mock
4import pytest
5import requests
7from superset_io.session import SupersetApiSession
10@pytest.fixture
11def fake_jwt():
12 header = (
13 base64.urlsafe_b64encode(b'{"alg":"HS256","typ":"JWT"}')
14 .decode("ascii")
15 .rstrip("=")
16 )
17 return f"{header}.payload.sig"
20class TestSupersetApiSession:
21 """Unit tests for SupersetApiSession."""
23 def test_init_sets_headers_conditionally(self):
24 s1 = SupersetApiSession(base_url="http://localhost:8088")
25 assert s1.base_url == "http://localhost:8088"
26 assert "User-Agent" in s1.headers
27 assert "Authorization" not in s1.headers
28 assert "X-CSRFToken" not in s1.headers
30 s2 = SupersetApiSession(
31 base_url="http://localhost:8088",
32 bearer_token="tok",
33 csrf_token="csrf",
34 )
35 assert s2.headers["Authorization"] == "Bearer tok"
36 assert s2.headers["X-CSRFToken"] == "csrf"
38 @pytest.mark.parametrize(
39 ("given", "expected"),
40 [
41 ("/api/v1/dashboard/", "http://localhost:8088/api/v1/dashboard/"),
42 ("https://example.com/api/test", "https://example.com/api/test"),
43 ],
44 )
45 def test_session_request_normalizes_urls(self, monkeypatch, given, expected):
46 session = SupersetApiSession(base_url="http://localhost:8088")
48 request_mock = Mock(return_value=Mock(status_code=200))
49 monkeypatch.setattr(requests.Session, "request", request_mock)
51 session.request("GET", given)
53 assert request_mock.call_args[0][0] == "GET"
54 assert request_mock.call_args[0][1] == expected
56 def test_from_token_http_error_does_not_crash_or_set_csrf(self, monkeypatch):
57 monkeypatch.setattr(
58 SupersetApiSession,
59 "_get_csrf_via_bearer",
60 Mock(side_effect=requests.HTTPError("no csrf")),
61 )
63 session = SupersetApiSession.from_token(
64 base_url="http://localhost:8088",
65 bearer_token="bearer-123",
66 )
68 assert session.bearer_token == "bearer-123"
69 assert session.headers["Authorization"] == "Bearer bearer-123"
70 assert session.csrf_token is None
71 assert "X-CSRFToken" not in session.headers
73 def test_from_token_success(self, monkeypatch):
74 monkeypatch.setattr(
75 SupersetApiSession,
76 "_get_csrf_via_bearer",
77 Mock(return_value="csrf-from-bearer"),
78 )
80 session = SupersetApiSession.from_token(
81 base_url="http://localhost:8088",
82 bearer_token="bearer-456",
83 )
85 assert session.bearer_token == "bearer-456"
86 assert session.headers["Authorization"] == "Bearer bearer-456"
87 assert session.csrf_token == "csrf-from-bearer"
88 assert session.headers["X-CSRFToken"] == "csrf-from-bearer"
90 def test_from_token_with_existing_session(self, monkeypatch):
91 existing_session = SupersetApiSession(base_url="http://localhost:8088")
93 monkeypatch.setattr(
94 SupersetApiSession,
95 "_get_csrf_via_bearer",
96 Mock(return_value="csrf-789"),
97 )
99 result = SupersetApiSession.from_token(
100 base_url="http://localhost:8088",
101 bearer_token="bearer-789",
102 session=existing_session,
103 )
105 assert result is existing_session
106 assert result.csrf_token == "csrf-789"
107 assert result.headers["X-CSRFToken"] == "csrf-789"
109 def test_get_csrf_via_session_cookie_no_csrf_field(self, monkeypatch):
110 session = SupersetApiSession(base_url="http://localhost:8088")
112 login_res = Mock()
113 login_res.text = "<html><body>No CSRF here</body></html>"
114 login_res.raise_for_status = Mock()
116 monkeypatch.setattr(session, "get", Mock(return_value=login_res))
118 with pytest.raises(RuntimeError, match="Could not find csrf_token field"):
119 session._get_csrf_via_session_cookie("admin", "password")
122class TestFromCredentials:
123 def test_from_credentials(self, monkeypatch):
124 monkeypatch.setattr(
125 SupersetApiSession,
126 "_get_bearer_token",
127 Mock(return_value="bearer-123"),
128 )
129 monkeypatch.setattr(
130 SupersetApiSession,
131 "_get_csrf_via_bearer",
132 Mock(return_value="csrf-abc"),
133 )
135 session = SupersetApiSession.from_credentials(
136 base_url="http://localhost:8088",
137 username="admin",
138 password="admin",
139 )
141 assert session.bearer_token == "bearer-123"
142 assert session.headers["Authorization"] == "Bearer bearer-123"
143 assert session.csrf_token == "csrf-abc"
144 assert session.headers["X-CSRFToken"] == "csrf-abc"
146 def test_from_credentials_fallback(self, monkeypatch):
147 monkeypatch.setattr(
148 SupersetApiSession,
149 "_get_bearer_token",
150 Mock(return_value="bearer-123"),
151 )
153 # Force the `from_token()` path to raise your custom RuntimeError
154 monkeypatch.setattr(
155 SupersetApiSession,
156 "_get_csrf_via_bearer",
157 Mock(side_effect=RuntimeError("Superset rejected the Bearer JWT")),
158 )
160 cookie_fallback = Mock()
161 monkeypatch.setattr(
162 SupersetApiSession,
163 "_get_csrf_via_session_cookie",
164 cookie_fallback,
165 )
167 SupersetApiSession.from_credentials(
168 base_url="http://localhost:8088",
169 username="admin",
170 password="admin",
171 )
173 cookie_fallback.assert_called_once()
175 def test_from_credentials_http_error(self, monkeypatch):
176 # Force bearer retrieval to fail
177 monkeypatch.setattr(
178 SupersetApiSession,
179 "_get_bearer_token",
180 Mock(side_effect=requests.HTTPError("login failed")),
181 )
183 # Patch from_token so we don't depend on CSRF internals here.
184 # Return the provided session unchanged.
185 from_token_mock = Mock(
186 side_effect=lambda base_url, bearer_token, session=None: session
187 )
188 monkeypatch.setattr(SupersetApiSession, "from_token", from_token_mock)
190 session = SupersetApiSession.from_credentials(
191 base_url="http://localhost:8088",
192 username="admin",
193 password="admin",
194 )
196 # No bearer token set, so no Authorization header should be present
197 assert session.bearer_token is None
198 assert "Authorization" not in session.headers
200 # from_token should still be called (current behavior), passing None as
201 # bearer token
202 from_token_mock.assert_called_once()
203 args, kwargs = from_token_mock.call_args
204 assert args[0] == "http://localhost:8088"
205 assert args[1] is None
206 assert kwargs["session"] is session
209class TestBearerToken:
210 def test_get_bearer_token_success(self, monkeypatch, fake_jwt):
211 session = SupersetApiSession(base_url="http://localhost:8088")
213 res = Mock()
214 res.json = Mock(return_value={"access_token": fake_jwt})
215 res.raise_for_status = Mock()
217 monkeypatch.setattr(session, "post", Mock(return_value=res))
219 token = session._get_bearer_token("admin", "password")
221 assert token == fake_jwt
223 def test_get_bearer_token_http_error(self, monkeypatch):
224 session = SupersetApiSession(base_url="http://localhost:8088")
226 res = Mock()
227 res.raise_for_status = Mock(side_effect=requests.HTTPError("Network error"))
229 monkeypatch.setattr(session, "post", Mock(return_value=res))
231 with pytest.raises(requests.HTTPError):
232 session._get_bearer_token("admin", "password")
234 def test_get_csrf_via_bearer(self, monkeypatch, fake_jwt):
235 session = SupersetApiSession(
236 base_url="http://localhost:8088",
237 bearer_token=fake_jwt,
238 )
240 res = Mock()
241 res.ok = True
242 res.json = Mock(return_value={"result": "csrf-abc"})
244 monkeypatch.setattr(session, "get", Mock(return_value=res))
246 assert session._get_csrf_via_bearer() == "csrf-abc"
248 def test_get_csrf_via_bearer_http_error(self, monkeypatch, fake_jwt):
249 session = SupersetApiSession(
250 base_url="http://localhost:8088",
251 bearer_token=fake_jwt,
252 )
254 res = Mock()
255 res.ok = False
256 res.status_code = 422
257 res.text = "The specified alg value is not allowed"
258 res.raise_for_status = Mock(side_effect=requests.HTTPError("login failed"))
260 monkeypatch.setattr(session, "get", Mock(return_value=res))
262 with pytest.raises(RuntimeError, match="Superset rejected the Bearer JWT"):
263 session._get_csrf_via_bearer()
266class TestSessionCookie:
267 def test_get_csrf_via_session_cookie_success(self, monkeypatch):
268 session = SupersetApiSession(base_url="http://localhost:8088")
270 # 1) GET /login/ returns HTML page
271 login_get_res = Mock()
272 login_get_res.text = '<input name="csrf_token" type="hidden" value="csrf-abc">'
273 login_get_res.raise_for_status = Mock()
275 # 2) POST /login/ succeeds (sets cookies in real life; here just no error)
276 login_post_res = Mock()
277 login_post_res.raise_for_status = Mock()
279 # 3) GET csrf endpoint returns JSON csrf
280 csrf_get_res = Mock()
281 csrf_get_res.raise_for_status = Mock()
282 csrf_get_res.json = Mock(return_value={"result": "csrf-abc"})
284 # Wire the call sequence
285 get_mock = Mock(side_effect=[login_get_res, csrf_get_res])
286 post_mock = Mock(return_value=login_post_res)
287 monkeypatch.setattr(session, "get", get_mock)
288 monkeypatch.setattr(session, "post", post_mock)
290 # Act
291 session._get_csrf_via_session_cookie("admin", "password")
293 # Assert: calls
294 get_mock.assert_any_call("/login/")
295 post_mock.assert_called_once_with(
296 "/login/",
297 data={
298 "username": "admin",
299 "password": "password",
300 "csrf_token": "csrf-abc",
301 },
302 allow_redirects=True,
303 headers={"Accept": "text/html,application/xhtml+xml"},
304 )
305 get_mock.assert_any_call("/api/v1/security/csrf_token/")
307 # Assert: session mutated
308 assert session.csrf_token == "csrf-abc"
309 assert session.headers["X-CSRFToken"] == "csrf-abc"