def test_graphql_uses_access_token_after_login(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post(
{
"data": {
"tenant": [{"id": tenant_id}],
"switch_tenant": {
"access_token": "ACCESS_TOKEN",
"expires_at": "2100-01-01",
"refresh_token": "REFRESH_TOKEN",
},
}
}
)
client = Client(api_token="api")
client.graphql({})
assert client.get_auth_token() == "api"
assert post.call_args[1]["headers"] == {
"Authorization": "Bearer api",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
client.login_to_tenant(tenant_id=tenant_id)
client.graphql({})
assert client.get_auth_token() == "ACCESS_TOKEN"
assert post.call_args[1]["headers"] == {
"Authorization": "Bearer ACCESS_TOKEN",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
def test_get_auth_token_refreshes_if_refresh_token_and_no_expiration(
self, monkeypatch):
refresh_token = MagicMock()
monkeypatch.setattr("prefect.Client._refresh_access_token",
refresh_token)
client = Client(api_token="api")
client._access_token = "access"
client._refresh_token = "refresh"
client._access_token_expires_at = None
client.get_auth_token()
assert refresh_token.called
def test_get_auth_token_refreshes_if_refresh_token_and_expiration_within_30_seconds(
self, monkeypatch):
refresh_token = MagicMock()
monkeypatch.setattr("prefect.Client._refresh_access_token",
refresh_token)
client = Client(api_token="api")
client._access_token = "access"
client._refresh_token = "refresh"
client._access_token_expires_at = pendulum.now().add(seconds=29)
client.get_auth_token()
assert refresh_token.called
def test_get_auth_token_doesnt_refresh_if_refresh_token_and_future_expiration(
self, monkeypatch):
refresh_token = MagicMock()
monkeypatch.setattr("prefect.Client._refresh_access_token",
refresh_token)
client = Client(api_token="api")
client._access_token = "access"
client._refresh_token = "refresh"
client._access_token_expires_at = pendulum.now().add(minutes=10)
assert client.get_auth_token() == "access"
refresh_token.assert_not_called()
def test_graphql_uses_access_token_after_login(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post(
{
"data": {
"tenant": [{"id": tenant_id}],
"switchTenant": {
"accessToken": "ACCESS_TOKEN",
"expiresIn": 600,
"refreshToken": "REFRESH_TOKEN",
},
}
}
)
client = Client(api_token="api")
client.graphql({})
assert client.get_auth_token() == "api"
assert post.call_args[1]["headers"] == dict(Authorization="Bearer api")
client.login_to_tenant(tenant_id=tenant_id)
client.graphql({})
assert client.get_auth_token() == "ACCESS_TOKEN"
assert post.call_args[1]["headers"] == dict(Authorization="Bearer ACCESS_TOKEN")
def test_get_auth_token_returns_access_token_if_set(self):
client = Client(api_token="api")
client._access_token = "access"
assert client.get_auth_token() == "access"
def test_get_auth_token_returns_api_if_access_token_not_set(self):
client = Client(api_token="api")
assert client._access_token is None
assert client.get_auth_token() == "api"
