Hugging Face's logo

Spaces:
jun2024
/
fc
Paused

App Files Community
fc
File size: 5,905 Bytes
b4f2238
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
"""Unit tests for McpApiKeyPool."""

from __future__ import annotations

from unittest.mock import patch

import pytest

from apps.enterprise_data.services.clients.api_key_pool import McpApiKeyPool


def _pool(*keys: str, provider: str = "test") -> McpApiKeyPool:
    return McpApiKeyPool(provider_name=provider, api_keys=list(keys))


# ── __init__ ──────────────────────────────────────────────────────────────────

def test_init_deduplicates_keys() -> None:
    pool = _pool("key1", "key1", "key2")
    assert pool.size == 2


def test_init_filters_empty_values() -> None:
    pool = _pool("key1", "", "  ", "key2")
    assert pool.size == 2


def test_init_empty_keys() -> None:
    pool = _pool()
    assert pool.size == 0


# ── ordered_keys ──────────────────────────────────────────────────────────────

def test_ordered_keys_single_key() -> None:
    pool = _pool("key1")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        mock_cache.get.return_value = None
        result = pool.ordered_keys()
    assert result == ["key1"]


def test_ordered_keys_multiple_keys_no_preferred() -> None:
    pool = _pool("key1", "key2", "key3")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        mock_cache.get.return_value = None
        result = pool.ordered_keys()
    assert set(result) == {"key1", "key2", "key3"}
    assert len(result) == 3


def test_ordered_keys_preferred_key_comes_first() -> None:
    pool = _pool("key1", "key2", "key3")
    preferred_fp = pool.fingerprint("key2")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        mock_cache.get.return_value = preferred_fp
        result = pool.ordered_keys()
    assert result[0] == "key2"


def test_ordered_keys_blocked_key_goes_last() -> None:
    pool = _pool("key1", "key2")
    blocked_fp = pool.fingerprint("key1")

    def cache_get(key: str) -> object:
        if "blocked" in key and blocked_fp in key:
            return True
        return None

    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        mock_cache.get.side_effect = cache_get
        result = pool.ordered_keys()

    assert result[-1] == "key1"
    assert result[0] == "key2"


# ── mark_success ──────────────────────────────────────────────────────────────

def test_mark_success_sets_preferred_and_clears_block() -> None:
    pool = _pool("key1", "key2")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        pool.mark_success("key1")
    mock_cache.set.assert_called_once()
    mock_cache.delete.assert_called_once()


def test_mark_success_empty_key_does_nothing() -> None:
    pool = _pool("key1")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        pool.mark_success("")
    mock_cache.set.assert_not_called()


# ── mark_auth_failed / mark_rate_limited ─────────────────────────────────────

def test_mark_auth_failed_sets_block() -> None:
    pool = _pool("key1")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        pool.mark_auth_failed("key1")
    mock_cache.set.assert_called_once()
    args = mock_cache.set.call_args
    # auth block TTL is 1 hour = 3600s
    assert args[1]["timeout"] == 3600


def test_mark_rate_limited_sets_block_with_short_ttl() -> None:
    pool = _pool("key1")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        pool.mark_rate_limited("key1")
    mock_cache.set.assert_called_once()
    args = mock_cache.set.call_args
    # rate limit block TTL is 2 minutes = 120s
    assert args[1]["timeout"] == 120


# ── fingerprint ───────────────────────────────────────────────────────────────

def test_fingerprint_returns_16_char_hex() -> None:
    pool = _pool("key1")
    fp = pool.fingerprint("key1")
    assert len(fp) == 16
    assert all(c in "0123456789abcdef" for c in fp)


def test_fingerprint_empty_returns_empty_string() -> None:
    pool = _pool()
    assert pool.fingerprint("") == ""
    assert pool.fingerprint("  ") == ""


def test_fingerprint_same_key_same_result() -> None:
    pool = _pool("key1")
    assert pool.fingerprint("key1") == pool.fingerprint("key1")


def test_fingerprint_different_keys_different_results() -> None:
    pool = _pool("key1", "key2")
    assert pool.fingerprint("key1") != pool.fingerprint("key2")


# ── θ‘₯ε……θΎΉηΌ˜εˆ†ζ”― ──────────────────────────────────────────────────────────────

def test_block_with_zero_ttl_does_nothing() -> None:
    # _block ttl_seconds <= 0 不写 cacheοΌˆθ¦†η›– line 71οΌ‰
    pool = _pool("key1")
    with patch("apps.enterprise_data.services.clients.api_key_pool.cache") as mock_cache:
        pool._block("key1", ttl_seconds=0)
    mock_cache.set.assert_not_called()


def test_order_with_preferred_preferred_not_in_list() -> None:
    # preferred fingerprint 不在 keys ι‡ŒοΌŒθΏ”ε›žεŽŸι‘ΊεΊοΌˆθ¦†η›– line 90οΌ‰
    pool = _pool("key1", "key2")
    result = pool._order_with_preferred(["key1", "key2"], "nonexistent_fp")
    assert result == ["key1", "key2"]