mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 12:39:27 +08:00
33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
|
from common.aimlapi_utils import DEFAULT_PARTNER_ID, DEFAULT_SOURCE, attribution_headers, partner_id, source
|
||
|
|
|
||
|
|
|
||
|
|
def test_defaults(monkeypatch):
|
||
|
|
monkeypatch.delenv("AIMLAPI_SOURCE", raising=False)
|
||
|
|
monkeypatch.delenv("AIMLAPI_PARTNER_ID", raising=False)
|
||
|
|
|
||
|
|
assert source() == DEFAULT_SOURCE
|
||
|
|
assert attribution_headers() == {
|
||
|
|
"X-AIMLAPI-Source": DEFAULT_SOURCE,
|
||
|
|
"X-AIMLAPI-Partner-ID": DEFAULT_PARTNER_ID,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def test_environment_overrides(monkeypatch):
|
||
|
|
monkeypatch.setenv("AIMLAPI_SOURCE", " agent/custom ")
|
||
|
|
monkeypatch.setenv("AIMLAPI_PARTNER_ID", " part_custom ")
|
||
|
|
|
||
|
|
assert source() == "agent/custom"
|
||
|
|
assert partner_id() == "part_custom"
|
||
|
|
assert attribution_headers() == {
|
||
|
|
"X-AIMLAPI-Source": "agent/custom",
|
||
|
|
"X-AIMLAPI-Partner-ID": "part_custom",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def test_blank_partner_id_is_omitted(monkeypatch):
|
||
|
|
monkeypatch.delenv("AIMLAPI_SOURCE", raising=False)
|
||
|
|
monkeypatch.setenv("AIMLAPI_PARTNER_ID", " ")
|
||
|
|
|
||
|
|
headers = attribution_headers()
|
||
|
|
assert headers == {"X-AIMLAPI-Source": DEFAULT_SOURCE}
|