mirror of
https://github.com/TencentCloud/TencentDB-Agent-Memory
synced 2026-07-11 21:04:29 +00:00
247 lines
8.7 KiB
Python
247 lines
8.7 KiB
Python
|
"""Unit tests for memory file reader module (cos.py)."""
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import hmac
|
||
|
|
import time
|
||
|
|
from typing import Any, Dict, Optional
|
||
|
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from tencentdb_agent_memory.cos import (
|
||
|
|
MemoryFileReader,
|
||
|
|
StsCredential,
|
||
|
|
StsCredentialManager,
|
||
|
|
_cos_v5_sign,
|
||
|
|
_parse_cos_url,
|
||
|
|
)
|
||
|
|
from tencentdb_agent_memory.errors import TDAMError
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# _parse_cos_url tests
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TestParseCosUrl:
|
||
|
|
def test_standard(self):
|
||
|
|
bucket, region = _parse_cos_url("https://my-test-bucket-1250000000.cos.ap-guangzhou.myqcloud.com")
|
||
|
|
assert bucket == "my-test-bucket-1250000000"
|
||
|
|
assert region == "ap-guangzhou"
|
||
|
|
|
||
|
|
def test_invalid_url(self):
|
||
|
|
with pytest.raises(TDAMError):
|
||
|
|
_parse_cos_url("https://invalid.example.com")
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# StsCredential tests
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TestStsCredential:
|
||
|
|
def test_parse(self):
|
||
|
|
cred = StsCredential({
|
||
|
|
"CosUrl": "https://my-bucket.cos.ap-guangzhou.myqcloud.com",
|
||
|
|
"TmpSecretId": "AK_test",
|
||
|
|
"TmpSecretKey": "SK_test",
|
||
|
|
"TmpToken": "tok",
|
||
|
|
"ExpirationTime": "2099-01-01T00:00:00Z",
|
||
|
|
"PathPrefix": "memory_v2/cos_data/mem-xxx",
|
||
|
|
})
|
||
|
|
assert cred.tmp_secret_id == "AK_test"
|
||
|
|
assert cred.bucket == "my-bucket"
|
||
|
|
assert cred.region == "ap-guangzhou"
|
||
|
|
assert cred.prefix == "memory_v2/cos_data/mem-xxx/"
|
||
|
|
assert cred.cos_host == "my-bucket.cos.ap-guangzhou.myqcloud.com"
|
||
|
|
assert cred.is_valid()
|
||
|
|
|
||
|
|
def test_prefix_trailing_slash(self):
|
||
|
|
cred = StsCredential({
|
||
|
|
"CosUrl": "https://b.cos.r.myqcloud.com",
|
||
|
|
"TmpSecretId": "AK",
|
||
|
|
"TmpSecretKey": "SK",
|
||
|
|
"TmpToken": "",
|
||
|
|
"ExpirationTime": "2099-01-01T00:00:00Z",
|
||
|
|
"PathPrefix": "pfx/",
|
||
|
|
})
|
||
|
|
assert cred.prefix == "pfx/"
|
||
|
|
|
||
|
|
def test_expired(self):
|
||
|
|
cred = StsCredential({
|
||
|
|
"CosUrl": "https://b.cos.r.myqcloud.com",
|
||
|
|
"TmpSecretId": "AK",
|
||
|
|
"TmpSecretKey": "SK",
|
||
|
|
"TmpToken": "",
|
||
|
|
"ExpirationTime": "2020-01-01T00:00:00Z",
|
||
|
|
"PathPrefix": "",
|
||
|
|
})
|
||
|
|
assert not cred.is_valid()
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# StsCredentialManager tests
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TestStsCredentialManager:
|
||
|
|
def _platform_response(self, expires_delta: float = 1800) -> dict:
|
||
|
|
from datetime import datetime, timezone, timedelta
|
||
|
|
exp = datetime.now(timezone.utc) + timedelta(seconds=expires_delta)
|
||
|
|
return {
|
||
|
|
"CosUrl": "https://test-bucket.cos.ap-guangzhou.myqcloud.com",
|
||
|
|
"TmpSecretId": "AK_test",
|
||
|
|
"TmpSecretKey": "SK_test",
|
||
|
|
"TmpToken": "tok_test",
|
||
|
|
"ExpirationTime": exp.isoformat(),
|
||
|
|
"PathPrefix": "test/",
|
||
|
|
}
|
||
|
|
|
||
|
|
@patch("tencentdb_agent_memory.cos.httpx.Client")
|
||
|
|
def test_fetch_on_first_call(self, MockClient):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 200
|
||
|
|
mock_resp.raise_for_status = MagicMock()
|
||
|
|
mock_resp.json.return_value = self._platform_response()
|
||
|
|
MockClient.return_value.post.return_value = mock_resp
|
||
|
|
|
||
|
|
mgr = StsCredentialManager(
|
||
|
|
endpoint="https://api.example.com",
|
||
|
|
api_key="sk-test",
|
||
|
|
service_id="mem-001",
|
||
|
|
)
|
||
|
|
cred = mgr.get_credential()
|
||
|
|
assert cred.tmp_secret_id == "AK_test"
|
||
|
|
MockClient.return_value.post.assert_called_once()
|
||
|
|
call_args = MockClient.return_value.post.call_args
|
||
|
|
assert "/v2/cos/secret" in call_args[0][0]
|
||
|
|
|
||
|
|
@patch("tencentdb_agent_memory.cos.httpx.Client")
|
||
|
|
def test_cache_hit(self, MockClient):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 200
|
||
|
|
mock_resp.raise_for_status = MagicMock()
|
||
|
|
mock_resp.json.return_value = self._platform_response()
|
||
|
|
MockClient.return_value.post.return_value = mock_resp
|
||
|
|
|
||
|
|
mgr = StsCredentialManager(
|
||
|
|
endpoint="https://api.example.com",
|
||
|
|
api_key="sk-test",
|
||
|
|
service_id="mem-001",
|
||
|
|
)
|
||
|
|
mgr.get_credential()
|
||
|
|
mgr.get_credential()
|
||
|
|
assert MockClient.return_value.post.call_count == 1
|
||
|
|
|
||
|
|
@patch("tencentdb_agent_memory.cos.httpx.Client")
|
||
|
|
def test_invalidate_forces_refetch(self, MockClient):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 200
|
||
|
|
mock_resp.raise_for_status = MagicMock()
|
||
|
|
mock_resp.json.return_value = self._platform_response()
|
||
|
|
MockClient.return_value.post.return_value = mock_resp
|
||
|
|
|
||
|
|
mgr = StsCredentialManager(
|
||
|
|
endpoint="https://api.example.com",
|
||
|
|
api_key="sk-test",
|
||
|
|
service_id="mem-001",
|
||
|
|
)
|
||
|
|
mgr.get_credential()
|
||
|
|
mgr.invalidate()
|
||
|
|
mgr.get_credential()
|
||
|
|
assert MockClient.return_value.post.call_count == 2
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# COS V5 signature tests
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TestCosV5Sign:
|
||
|
|
def test_signature_format(self):
|
||
|
|
auth = _cos_v5_sign(
|
||
|
|
secret_id="AKID_test",
|
||
|
|
secret_key="SK_test",
|
||
|
|
method="GET",
|
||
|
|
path="/test/file.md",
|
||
|
|
host="bucket.cos.ap-guangzhou.myqcloud.com",
|
||
|
|
start_time=1000000,
|
||
|
|
end_time=1000600,
|
||
|
|
)
|
||
|
|
assert "q-sign-algorithm=sha1" in auth
|
||
|
|
assert "q-ak=AKID_test" in auth
|
||
|
|
assert "q-sign-time=1000000;1000600" in auth
|
||
|
|
assert "q-signature=" in auth
|
||
|
|
|
||
|
|
def test_deterministic(self):
|
||
|
|
kwargs = dict(
|
||
|
|
secret_id="AK", secret_key="SK",
|
||
|
|
method="GET", path="/a.md",
|
||
|
|
host="b.cos.r.myqcloud.com",
|
||
|
|
start_time=100, end_time=200,
|
||
|
|
)
|
||
|
|
assert _cos_v5_sign(**kwargs) == _cos_v5_sign(**kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# MemoryFileReader tests (mock HTTP)
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TestMemoryFileReader:
|
||
|
|
def _make_reader(self, http_response: Any = None) -> tuple:
|
||
|
|
from datetime import datetime, timezone, timedelta
|
||
|
|
exp = datetime.now(timezone.utc) + timedelta(seconds=1800)
|
||
|
|
platform_resp = {
|
||
|
|
"CosUrl": "https://bkt.cos.ap-gz.myqcloud.com",
|
||
|
|
"TmpSecretId": "AK",
|
||
|
|
"TmpSecretKey": "SK",
|
||
|
|
"TmpToken": "tok",
|
||
|
|
"ExpirationTime": exp.isoformat(),
|
||
|
|
"PathPrefix": "pfx/",
|
||
|
|
}
|
||
|
|
mock_sts_resp = MagicMock()
|
||
|
|
mock_sts_resp.status_code = 200
|
||
|
|
mock_sts_resp.raise_for_status = MagicMock()
|
||
|
|
mock_sts_resp.json.return_value = platform_resp
|
||
|
|
|
||
|
|
with patch("tencentdb_agent_memory.cos.httpx.Client") as MockClient:
|
||
|
|
MockClient.return_value.post.return_value = mock_sts_resp
|
||
|
|
mgr = StsCredentialManager(
|
||
|
|
endpoint="https://api.example.com",
|
||
|
|
api_key="sk-test",
|
||
|
|
service_id="mem-001",
|
||
|
|
)
|
||
|
|
# Pre-populate credential
|
||
|
|
mgr.get_credential()
|
||
|
|
|
||
|
|
mock_cos_client = MagicMock()
|
||
|
|
if http_response is not None:
|
||
|
|
mock_cos_client.get.return_value = http_response
|
||
|
|
reader = MemoryFileReader(mgr, client=mock_cos_client)
|
||
|
|
return reader, mock_cos_client
|
||
|
|
|
||
|
|
def test_read_success(self):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 200
|
||
|
|
mock_resp.text = "# Hello World"
|
||
|
|
reader, client = self._make_reader(mock_resp)
|
||
|
|
result = reader.read("scene_blocks/test.md")
|
||
|
|
assert result == "# Hello World"
|
||
|
|
client.get.assert_called_once()
|
||
|
|
call_url = client.get.call_args[0][0]
|
||
|
|
assert "pfx/scene_blocks/test.md" in call_url
|
||
|
|
|
||
|
|
def test_read_404(self):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 404
|
||
|
|
reader, _ = self._make_reader(mock_resp)
|
||
|
|
with pytest.raises(TDAMError) as exc_info:
|
||
|
|
reader.read("nonexistent.md")
|
||
|
|
assert exc_info.value.code == 404
|
||
|
|
|
||
|
|
def test_security_token_header(self):
|
||
|
|
mock_resp = MagicMock()
|
||
|
|
mock_resp.status_code = 200
|
||
|
|
mock_resp.text = "ok"
|
||
|
|
reader, client = self._make_reader(mock_resp)
|
||
|
|
reader.read("test.md")
|
||
|
|
call_headers = client.get.call_args[1]["headers"]
|
||
|
|
assert "x-cos-security-token" in call_headers
|
||
|
|
assert call_headers["x-cos-security-token"] == "tok"
|