270 lines
9.4 KiB
Python
270 lines
9.4 KiB
Python
"""Тесты переиспользуемого модуля upload (backend): session, safe_name, blueprint.
|
|
|
|
Запуск:
|
|
python3 tests/test_upload_module.py
|
|
pytest tests/test_upload_module.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# Корень проекта — для импорта пакета upload
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from flask import Flask
|
|
|
|
from upload.backend.session import (
|
|
create_session, add_file, get_files, file_count,
|
|
store_result, get_result, store_csv, get_csv,
|
|
cleanup, request_cancel, get_cancel_event,
|
|
)
|
|
from upload.backend.upload_refs import safe_name, create_upload_refs_blueprint
|
|
|
|
VM_PREFIX = "https://contracts.kube5s.ru/drhider-upload/"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# session
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
def test_create_session():
|
|
sid = create_session()
|
|
assert len(sid) == 32 # uuid4().hex
|
|
|
|
|
|
def test_add_file_and_get():
|
|
sid = create_session()
|
|
assert add_file(sid, "a.txt", b"hello") is True
|
|
assert get_files(sid) == [("a.txt", b"hello")]
|
|
assert file_count(sid) == 1
|
|
|
|
|
|
def test_add_file_missing_session():
|
|
assert add_file("nonexistent", "a.txt", b"x") is False
|
|
assert get_files("nonexistent") is None
|
|
|
|
|
|
def test_session_limit(monkeypatch):
|
|
# importlib: `import ...session.add_file as m` затеняется одноимённой функцией
|
|
import importlib
|
|
add_file_mod = importlib.import_module("upload.backend.session.add_file")
|
|
old = add_file_mod.MAX_SESSION_BYTES
|
|
add_file_mod.MAX_SESSION_BYTES = 10
|
|
try:
|
|
sid = create_session()
|
|
assert add_file(sid, "a", b"12345") is True
|
|
# суммарно уже 5; +6 = 11 > 10 → False
|
|
assert add_file(sid, "b", b"123456") is False
|
|
assert file_count(sid) == 1
|
|
finally:
|
|
add_file_mod.MAX_SESSION_BYTES = old
|
|
|
|
|
|
def test_store_result_and_csv():
|
|
sid = create_session()
|
|
assert store_result(sid, b"zip") is True
|
|
assert get_result(sid) == b"zip"
|
|
assert store_csv(sid, "a,b") is True
|
|
assert get_csv(sid) == "a,b"
|
|
|
|
|
|
def test_cleanup():
|
|
sid = create_session()
|
|
cleanup(sid)
|
|
assert get_files(sid) is None
|
|
|
|
|
|
def test_cancel():
|
|
sid = create_session()
|
|
assert request_cancel(sid) is True
|
|
ev = get_cancel_event(sid)
|
|
assert ev is not None and ev.is_set()
|
|
assert request_cancel("nonexistent") is False
|
|
assert get_cancel_event("nonexistent") is None
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# safe_name
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
def test_safe_name():
|
|
assert safe_name("a.txt") == "a.txt"
|
|
assert safe_name("dir/sub/a.txt") == "dir/sub/a.txt"
|
|
assert safe_name("dir\\sub\\a.txt") == "dir/sub/a.txt"
|
|
assert safe_name("../etc/passwd") == ""
|
|
assert safe_name("a/../../b") == ""
|
|
assert safe_name("") == ""
|
|
assert safe_name("./a.txt") == "a.txt"
|
|
assert safe_name("/etc/passwd") == "etc/passwd"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# blueprint upload_refs
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
class _FakeStream:
|
|
def __init__(self, payload=b"", status=200, exc=None):
|
|
self._payload = payload
|
|
self._status = status
|
|
self._exc = exc
|
|
|
|
def __enter__(self):
|
|
if self._exc:
|
|
raise self._exc
|
|
return self
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
def raise_for_status(self):
|
|
if self._status >= 400:
|
|
raise RuntimeError("HTTP %d" % self._status)
|
|
|
|
def iter_bytes(self):
|
|
yield self._payload
|
|
|
|
|
|
class _FakeClient:
|
|
"""Имитация httpx.Client: stream/delete без сети.
|
|
|
|
behavior: dict url -> "ok" | {"payload": bytes} | {"exc": Exception}
|
|
"""
|
|
|
|
def __init__(self, behavior):
|
|
self._behavior = behavior
|
|
self.deleted = []
|
|
self.streams = []
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
def stream(self, method, url):
|
|
self.streams.append((method, url))
|
|
b = self._behavior.get(url)
|
|
if isinstance(b, dict) and "payload" in b:
|
|
return _FakeStream(payload=b["payload"])
|
|
if isinstance(b, dict) and "exc" in b:
|
|
return _FakeStream(exc=b["exc"])
|
|
if b == "ok":
|
|
return _FakeStream()
|
|
return _FakeStream(status=500)
|
|
|
|
def delete(self, url):
|
|
self.deleted.append(url)
|
|
|
|
|
|
def _make_client(behavior, cfg=None, monkeypatch=None):
|
|
import upload.backend.upload_refs.blueprint as bp_mod
|
|
monkeypatch.setattr(bp_mod.httpx, "Client", lambda **kw: _FakeClient(behavior))
|
|
app = Flask(__name__)
|
|
app.register_blueprint(create_upload_refs_blueprint(cfg or {}))
|
|
return app.test_client()
|
|
|
|
|
|
def test_upload_refs_no_files(monkeypatch):
|
|
c = _make_client({}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={"session": "", "files": []})
|
|
assert rv.status_code == 400
|
|
|
|
|
|
def test_upload_refs_ok(monkeypatch):
|
|
url = VM_PREFIX + "tok_0"
|
|
c = _make_client({url: {"payload": b"hello"}}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={
|
|
"session": "",
|
|
"files": [{"name": "a.txt", "size": 5, "url": url}],
|
|
})
|
|
assert rv.status_code == 200
|
|
data = rv.get_json()
|
|
assert data["ok"] is True
|
|
assert data["count"] == 1
|
|
files = get_files(data["session"])
|
|
assert files == [("a.txt", b"hello")]
|
|
|
|
|
|
def test_upload_refs_ssrf(monkeypatch):
|
|
c = _make_client({}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={
|
|
"session": "",
|
|
"files": [{"name": "evil.txt", "size": 5, "url": "https://evil.example/x"}],
|
|
})
|
|
assert rv.status_code == 200
|
|
data = rv.get_json()
|
|
assert data["ok"] is True
|
|
assert data["count"] == 0 # unsafe URL пропущен
|
|
|
|
|
|
def test_upload_refs_too_large(monkeypatch):
|
|
url = VM_PREFIX + "tok_0"
|
|
c = _make_client({}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={
|
|
"session": "",
|
|
"files": [{"name": "big.bin", "size": 60 * 1024 * 1024, "url": url}],
|
|
})
|
|
assert rv.status_code == 200
|
|
assert rv.get_json()["count"] == 0 # сверх 50МБ → skip+delete
|
|
|
|
|
|
def test_upload_refs_session_not_found(monkeypatch):
|
|
url = VM_PREFIX + "tok_0"
|
|
c = _make_client({url: {"payload": b"x"}}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={
|
|
"session": "no-such-session",
|
|
"files": [{"name": "a.txt", "size": 1, "url": url}],
|
|
})
|
|
assert rv.status_code == 404
|
|
assert rv.get_json()["error"] == "Session not found"
|
|
|
|
|
|
def test_upload_refs_pull_failed(monkeypatch):
|
|
url = VM_PREFIX + "tok_0"
|
|
behavior = {url: {"exc": RuntimeError("dns fail")}}
|
|
c = _make_client(behavior, cfg={"pullRetryDelay": 0}, monkeypatch=monkeypatch)
|
|
rv = c.post("/api/upload_refs", json={
|
|
"session": "",
|
|
"files": [{"name": "a.txt", "size": 1, "url": url}],
|
|
})
|
|
assert rv.status_code == 502 # все ретраи провалились
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# main (запуск без pytest)
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
|
|
def main():
|
|
"""Прогнать все test_* функции без pytest."""
|
|
import inspect
|
|
import traceback
|
|
fns = [(k, v) for k, v in sorted(globals().items()) if k.startswith("test_")]
|
|
mp = _FakeMonkeyPatch()
|
|
failed = 0
|
|
for name, fn in fns:
|
|
try:
|
|
if "monkeypatch" in inspect.signature(fn).parameters:
|
|
fn(monkeypatch=mp)
|
|
else:
|
|
fn()
|
|
print("PASS %s" % name)
|
|
except Exception:
|
|
failed += 1
|
|
print("FAIL %s" % name)
|
|
traceback.print_exc()
|
|
if failed:
|
|
print("%d тестов упало" % failed)
|
|
sys.exit(1)
|
|
print("Все тесты прошли")
|
|
|
|
|
|
class _FakeMonkeyPatch:
|
|
"""Мини-monkeypatch: setattr, работает до конца вызова (без undo)."""
|
|
|
|
def setattr(self, target, name, value):
|
|
setattr(target, name, value)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|