96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
# 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.
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from tests.integration_tests.test_app import app
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"base_url, expected_referer",
|
|
[
|
|
("http://base-url", None),
|
|
("http://base-url/", None),
|
|
("https://base-url", "https://base-url/api/v1/chart/warm_up_cache"),
|
|
("https://base-url/", "https://base-url/api/v1/chart/warm_up_cache"),
|
|
],
|
|
ids=[
|
|
"Without trailing slash (HTTP)",
|
|
"With trailing slash (HTTP)",
|
|
"Without trailing slash (HTTPS)",
|
|
"With trailing slash (HTTPS)",
|
|
],
|
|
)
|
|
@mock.patch("superset.tasks.cache.fetch_csrf_token")
|
|
@mock.patch("superset.tasks.cache.request.Request")
|
|
@mock.patch("superset.tasks.cache.request.urlopen")
|
|
@mock.patch("superset.tasks.cache.is_secure_url")
|
|
def test_fetch_url(
|
|
mock_is_secure_url,
|
|
mock_urlopen,
|
|
mock_request_cls,
|
|
mock_fetch_csrf_token,
|
|
base_url,
|
|
expected_referer,
|
|
):
|
|
from superset.tasks.cache import fetch_url
|
|
|
|
mock_request = mock.MagicMock()
|
|
mock_request_cls.return_value = mock_request
|
|
|
|
mock_urlopen.return_value = mock.MagicMock()
|
|
mock_urlopen.return_value.code = 200
|
|
|
|
# Mock the URL validation to return True for HTTPS and False for HTTP
|
|
mock_is_secure_url.return_value = base_url.startswith("https")
|
|
|
|
initial_headers = {"Cookie": "cookie", "key": "value"}
|
|
csrf_headers = initial_headers | {"X-CSRF-Token": "csrf_token"}
|
|
|
|
# Conditionally add the Referer header and assert its presence
|
|
if expected_referer:
|
|
csrf_headers = csrf_headers | {"Referer": expected_referer}
|
|
assert csrf_headers["Referer"] == expected_referer
|
|
|
|
mock_fetch_csrf_token.return_value = csrf_headers
|
|
|
|
app.config["WEBDRIVER_BASEURL"] = base_url
|
|
data = "data"
|
|
data_encoded = b"data"
|
|
|
|
result = fetch_url(data, initial_headers)
|
|
|
|
expected_url = (
|
|
f"{base_url}/api/v1/chart/warm_up_cache"
|
|
if not base_url.endswith("/")
|
|
else f"{base_url}api/v1/chart/warm_up_cache"
|
|
)
|
|
|
|
mock_fetch_csrf_token.assert_called_once_with(initial_headers)
|
|
|
|
mock_request_cls.assert_called_once_with(
|
|
expected_url, # Use the dynamic URL based on base_url
|
|
data=data_encoded,
|
|
headers=csrf_headers,
|
|
method="PUT",
|
|
)
|
|
# assert the same Request object is used
|
|
mock_urlopen.assert_called_once_with(mock_request, timeout=mock.ANY)
|
|
|
|
assert data == result["success"]
|