# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name, disallowed-name from datetime import datetime from freezegun import freeze_time from pytest_mock import MockerFixture from superset.utils.oauth2 import get_oauth2_access_token def test_get_oauth2_access_token_base_no_token(mocker: MockerFixture) -> None: """ Test `get_oauth2_access_token` when there's no token. """ db = mocker.patch("superset.utils.oauth2.db") db_engine_spec = mocker.MagicMock() db.session.query().filter_by().one_or_none.return_value = None assert get_oauth2_access_token({}, 1, 1, db_engine_spec) is None def test_get_oauth2_access_token_base_token_valid(mocker: MockerFixture) -> None: """ Test `get_oauth2_access_token` when the token is valid. """ db = mocker.patch("superset.utils.oauth2.db") db_engine_spec = mocker.MagicMock() token = mocker.MagicMock() token.access_token = "access-token" token.access_token_expiration = datetime(2024, 1, 2) db.session.query().filter_by().one_or_none.return_value = token with freeze_time("2024-01-01"): assert get_oauth2_access_token({}, 1, 1, db_engine_spec) == "access-token" def test_get_oauth2_access_token_base_refresh(mocker: MockerFixture) -> None: """ Test `get_oauth2_access_token` when the token needs to be refreshed. """ db = mocker.patch("superset.utils.oauth2.db") db_engine_spec = mocker.MagicMock() db_engine_spec.get_oauth2_fresh_token.return_value = { "access_token": "new-token", "expires_in": 3600, } token = mocker.MagicMock() token.access_token = "access-token" token.access_token_expiration = datetime(2024, 1, 1) token.refresh_token = "refresh-token" db.session.query().filter_by().one_or_none.return_value = token with freeze_time("2024-01-02"): assert get_oauth2_access_token({}, 1, 1, db_engine_spec) == "new-token" # check that token was updated assert token.access_token == "new-token" assert token.access_token_expiration == datetime(2024, 1, 2, 1) db.session.add.assert_called_with(token) def test_get_oauth2_access_token_base_no_refresh(mocker: MockerFixture) -> None: """ Test `get_oauth2_access_token` when token is expired and there's no refresh. """ db = mocker.patch("superset.utils.oauth2.db") db_engine_spec = mocker.MagicMock() token = mocker.MagicMock() token.access_token = "access-token" token.access_token_expiration = datetime(2024, 1, 1) token.refresh_token = None db.session.query().filter_by().one_or_none.return_value = token with freeze_time("2024-01-02"): assert get_oauth2_access_token({}, 1, 1, db_engine_spec) is None # check that token was deleted db.session.delete.assert_called_with(token)