feat: add archive storage client and env config (#30422)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-05-03 00:48:04 +08:00 · 2025-12-31 16:14:46 +08:00
parent 2bb1e24fb4
commit 3015e9be73
7 changed files with 689 additions and 0 deletions
--- a/api/tests/unit_tests/libs/test_archive_storage.py
+++ b/api/tests/unit_tests/libs/test_archive_storage.py
@ -0,0 +1,272 @@
+import base64
+import hashlib
+from datetime import datetime
+from unittest.mock import ANY, MagicMock
+
+import pytest
+from botocore.exceptions import ClientError
+
+from libs import archive_storage as storage_module
+from libs.archive_storage import (
+    ArchiveStorage,
+    ArchiveStorageError,
+    ArchiveStorageNotConfiguredError,
+)
+
+BUCKET_NAME = "archive-bucket"
+
+
+def _configure_storage(monkeypatch, **overrides):
+    defaults = {
+        "ARCHIVE_STORAGE_ENABLED": True,
+        "ARCHIVE_STORAGE_ENDPOINT": "https://storage.example.com",
+        "ARCHIVE_STORAGE_ARCHIVE_BUCKET": BUCKET_NAME,
+        "ARCHIVE_STORAGE_ACCESS_KEY": "access",
+        "ARCHIVE_STORAGE_SECRET_KEY": "secret",
+        "ARCHIVE_STORAGE_REGION": "auto",
+    }
+    defaults.update(overrides)
+    for key, value in defaults.items():
+        monkeypatch.setattr(storage_module.dify_config, key, value, raising=False)
+
+
+def _client_error(code: str) -> ClientError:
+    return ClientError({"Error": {"Code": code}}, "Operation")
+
+
+def _mock_client(monkeypatch):
+    client = MagicMock()
+    client.head_bucket.return_value = None
+    boto_client = MagicMock(return_value=client)
+    monkeypatch.setattr(storage_module.boto3, "client", boto_client)
+    return client, boto_client
+
+
+def test_init_disabled(monkeypatch):
+    _configure_storage(monkeypatch, ARCHIVE_STORAGE_ENABLED=False)
+    with pytest.raises(ArchiveStorageNotConfiguredError, match="not enabled"):
+        ArchiveStorage(bucket=BUCKET_NAME)
+
+
+def test_init_missing_config(monkeypatch):
+    _configure_storage(monkeypatch, ARCHIVE_STORAGE_ENDPOINT=None)
+    with pytest.raises(ArchiveStorageNotConfiguredError, match="incomplete"):
+        ArchiveStorage(bucket=BUCKET_NAME)
+
+
+def test_init_bucket_not_found(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.head_bucket.side_effect = _client_error("404")
+
+    with pytest.raises(ArchiveStorageNotConfiguredError, match="does not exist"):
+        ArchiveStorage(bucket=BUCKET_NAME)
+
+
+def test_init_bucket_access_denied(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.head_bucket.side_effect = _client_error("403")
+
+    with pytest.raises(ArchiveStorageNotConfiguredError, match="Access denied"):
+        ArchiveStorage(bucket=BUCKET_NAME)
+
+
+def test_init_bucket_other_error(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.head_bucket.side_effect = _client_error("500")
+
+    with pytest.raises(ArchiveStorageError, match="Failed to access archive bucket"):
+        ArchiveStorage(bucket=BUCKET_NAME)
+
+
+def test_init_sets_client(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, boto_client = _mock_client(monkeypatch)
+
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    boto_client.assert_called_once_with(
+        "s3",
+        endpoint_url="https://storage.example.com",
+        aws_access_key_id="access",
+        aws_secret_access_key="secret",
+        region_name="auto",
+        config=ANY,
+    )
+    assert storage.client is client
+    assert storage.bucket == BUCKET_NAME
+
+
+def test_put_object_returns_checksum(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    data = b"hello"
+    checksum = storage.put_object("key", data)
+
+    expected_md5 = hashlib.md5(data).hexdigest()
+    expected_content_md5 = base64.b64encode(hashlib.md5(data).digest()).decode()
+    client.put_object.assert_called_once_with(
+        Bucket="archive-bucket",
+        Key="key",
+        Body=data,
+        ContentMD5=expected_content_md5,
+    )
+    assert checksum == expected_md5
+
+
+def test_put_object_raises_on_error(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+    client.put_object.side_effect = _client_error("500")
+
+    with pytest.raises(ArchiveStorageError, match="Failed to upload object"):
+        storage.put_object("key", b"data")
+
+
+def test_get_object_returns_bytes(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    body = MagicMock()
+    body.read.return_value = b"payload"
+    client.get_object.return_value = {"Body": body}
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    assert storage.get_object("key") == b"payload"
+
+
+def test_get_object_missing(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.get_object.side_effect = _client_error("NoSuchKey")
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    with pytest.raises(FileNotFoundError, match="Archive object not found"):
+        storage.get_object("missing")
+
+
+def test_get_object_stream(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    body = MagicMock()
+    body.iter_chunks.return_value = [b"a", b"b"]
+    client.get_object.return_value = {"Body": body}
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    assert list(storage.get_object_stream("key")) == [b"a", b"b"]
+
+
+def test_get_object_stream_missing(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.get_object.side_effect = _client_error("NoSuchKey")
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    with pytest.raises(FileNotFoundError, match="Archive object not found"):
+        list(storage.get_object_stream("missing"))
+
+
+def test_object_exists(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    assert storage.object_exists("key") is True
+    client.head_object.side_effect = _client_error("404")
+    assert storage.object_exists("missing") is False
+
+
+def test_delete_object_error(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.delete_object.side_effect = _client_error("500")
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    with pytest.raises(ArchiveStorageError, match="Failed to delete object"):
+        storage.delete_object("key")
+
+
+def test_list_objects(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    paginator = MagicMock()
+    paginator.paginate.return_value = [
+        {"Contents": [{"Key": "a"}, {"Key": "b"}]},
+        {"Contents": [{"Key": "c"}]},
+    ]
+    client.get_paginator.return_value = paginator
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    assert storage.list_objects("prefix") == ["a", "b", "c"]
+    paginator.paginate.assert_called_once_with(Bucket="archive-bucket", Prefix="prefix")
+
+
+def test_list_objects_error(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    paginator = MagicMock()
+    paginator.paginate.side_effect = _client_error("500")
+    client.get_paginator.return_value = paginator
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    with pytest.raises(ArchiveStorageError, match="Failed to list objects"):
+        storage.list_objects("prefix")
+
+
+def test_generate_presigned_url(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.generate_presigned_url.return_value = "http://signed-url"
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    url = storage.generate_presigned_url("key", expires_in=123)
+
+    client.generate_presigned_url.assert_called_once_with(
+        ClientMethod="get_object",
+        Params={"Bucket": "archive-bucket", "Key": "key"},
+        ExpiresIn=123,
+    )
+    assert url == "http://signed-url"
+
+
+def test_generate_presigned_url_error(monkeypatch):
+    _configure_storage(monkeypatch)
+    client, _ = _mock_client(monkeypatch)
+    client.generate_presigned_url.side_effect = _client_error("500")
+    storage = ArchiveStorage(bucket=BUCKET_NAME)
+
+    with pytest.raises(ArchiveStorageError, match="Failed to generate pre-signed URL"):
+        storage.generate_presigned_url("key")
+
+
+def test_serialization_roundtrip():
+    records = [
+        {
+            "id": "1",
+            "created_at": datetime(2024, 1, 1, 12, 0, 0),
+            "payload": {"nested": "value"},
+            "items": [{"name": "a"}],
+        },
+        {"id": "2", "value": 123},
+    ]
+
+    data = ArchiveStorage.serialize_to_jsonl_gz(records)
+    decoded = ArchiveStorage.deserialize_from_jsonl_gz(data)
+
+    assert decoded[0]["id"] == "1"
+    assert decoded[0]["payload"]["nested"] == "value"
+    assert decoded[0]["items"][0]["name"] == "a"
+    assert "2024-01-01T12:00:00" in decoded[0]["created_at"]
+    assert decoded[1]["value"] == 123
+
+
+def test_content_md5_matches_checksum():
+    data = b"checksum"
+    expected = base64.b64encode(hashlib.md5(data).digest()).decode()
+
+    assert ArchiveStorage._content_md5(data) == expected
+    assert ArchiveStorage.compute_checksum(data) == hashlib.md5(data).hexdigest()