222 lines
7.2 KiB
Python
222 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for the Danbooru API.
|
|
|
|
Verifies:
|
|
1. Authentication with the API key works
|
|
2. Tag listing endpoint returns expected fields
|
|
3. Pagination works (multiple pages)
|
|
4. Tag search / filtering by category works
|
|
|
|
Usage:
|
|
python scripts/test_danbooru_api.py
|
|
|
|
Reads DANBOORU_API_KEY from .env or environment.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
try:
|
|
import requests
|
|
except ImportError:
|
|
print("ERROR: 'requests' is not installed. Run: pip install requests")
|
|
sys.exit(1)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Load .env
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_env() -> dict[str, str]:
|
|
env: dict[str, str] = {}
|
|
env_path = Path(__file__).parent.parent / ".env"
|
|
if env_path.exists():
|
|
for line in env_path.read_text().splitlines():
|
|
line = line.strip()
|
|
if line and not line.startswith("#") and "=" in line:
|
|
k, _, v = line.partition("=")
|
|
env[k.strip()] = v.strip()
|
|
return env
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
BASE_URL = "https://danbooru.donmai.us"
|
|
|
|
# Danbooru tag categories
|
|
CATEGORY_NAMES = {
|
|
0: "general",
|
|
1: "artist",
|
|
3: "copyright",
|
|
4: "character",
|
|
5: "meta",
|
|
}
|
|
|
|
|
|
def make_session(api_key: str | None = None, username: str | None = None) -> requests.Session:
|
|
"""Create a requests Session.
|
|
|
|
Danbooru public endpoints (tag listing, searching) do not require
|
|
authentication. Auth is only needed for account-specific actions.
|
|
When provided, credentials must be (login, api_key) — NOT (user, api_key).
|
|
"""
|
|
session = requests.Session()
|
|
if api_key and username:
|
|
session.auth = (username, api_key)
|
|
session.headers.update({"User-Agent": "danbooru-mcp-test/0.1"})
|
|
return session
|
|
|
|
|
|
def get_tags_page(
|
|
session: requests.Session,
|
|
page: int = 1,
|
|
limit: int = 20,
|
|
search_name: str | None = None,
|
|
search_category: int | None = None,
|
|
order: str = "count", # "count" | "name" | "date"
|
|
) -> list[dict]:
|
|
"""Fetch one page of tags from the Danbooru API."""
|
|
params: dict = {
|
|
"page": page,
|
|
"limit": limit,
|
|
"search[order]": order,
|
|
}
|
|
if search_name:
|
|
params["search[name_matches]"] = search_name
|
|
if search_category is not None:
|
|
params["search[category]"] = search_category
|
|
|
|
resp = session.get(f"{BASE_URL}/tags.json", params=params, timeout=15)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_basic_fetch(session: requests.Session) -> None:
|
|
print("\n[1] Basic fetch — top 5 tags by post count")
|
|
tags = get_tags_page(session, page=1, limit=5, order="count")
|
|
assert isinstance(tags, list), f"Expected list, got {type(tags)}"
|
|
assert len(tags) > 0, "No tags returned"
|
|
|
|
for tag in tags:
|
|
cat = CATEGORY_NAMES.get(tag.get("category", -1), "unknown")
|
|
print(f" [{cat:12s}] {tag['name']:40s} posts={tag['post_count']:>8,}")
|
|
|
|
print(" PASS ✓")
|
|
|
|
|
|
def test_fields_present(session: requests.Session) -> None:
|
|
print("\n[2] Field presence check")
|
|
tags = get_tags_page(session, page=1, limit=1, order="count")
|
|
tag = tags[0]
|
|
required = {"id", "name", "post_count", "category", "is_deprecated", "words"}
|
|
missing = required - set(tag.keys())
|
|
assert not missing, f"Missing fields: {missing}"
|
|
print(f" Fields present: {sorted(tag.keys())}")
|
|
print(f" Sample tag: name={tag['name']!r} category={CATEGORY_NAMES.get(tag['category'])} deprecated={tag['is_deprecated']}")
|
|
print(" PASS ✓")
|
|
|
|
|
|
def test_pagination(session: requests.Session) -> None:
|
|
print("\n[3] Pagination — page 1 vs page 2 should differ")
|
|
p1 = get_tags_page(session, page=1, limit=5, order="count")
|
|
time.sleep(0.5)
|
|
p2 = get_tags_page(session, page=2, limit=5, order="count")
|
|
names_p1 = {t["name"] for t in p1}
|
|
names_p2 = {t["name"] for t in p2}
|
|
overlap = names_p1 & names_p2
|
|
assert not overlap, f"Pages 1 and 2 share tags: {overlap}"
|
|
print(f" Page 1: {sorted(names_p1)}")
|
|
print(f" Page 2: {sorted(names_p2)}")
|
|
print(" PASS ✓")
|
|
|
|
|
|
def test_category_filter(session: requests.Session) -> None:
|
|
print("\n[4] Category filter — fetch only 'character' tags (category=4)")
|
|
tags = get_tags_page(session, page=1, limit=5, search_category=4, order="count")
|
|
for tag in tags:
|
|
assert tag["category"] == 4, f"Expected category 4, got {tag['category']} for {tag['name']}"
|
|
print(f" {tag['name']:40s} posts={tag['post_count']:>8,}")
|
|
print(" PASS ✓")
|
|
|
|
|
|
def test_name_search(session: requests.Session) -> None:
|
|
print("\n[5] Name search — tags matching 'blue_hair*'")
|
|
tags = get_tags_page(session, page=1, limit=5, search_name="blue_hair*", order="count")
|
|
assert len(tags) > 0, "No results for blue_hair*"
|
|
for tag in tags:
|
|
cat = CATEGORY_NAMES.get(tag.get("category", -1), "unknown")
|
|
print(f" [{cat:12s}] {tag['name']:40s} posts={tag['post_count']:>8,}")
|
|
print(" PASS ✓")
|
|
|
|
|
|
def test_well_known_tags(session: requests.Session) -> None:
|
|
print("\n[6] Well-known tags — '1girl', 'blue_hair', 'sword' should exist")
|
|
for tag_name in ("1girl", "blue_hair", "sword"):
|
|
tags = get_tags_page(session, page=1, limit=1, search_name=tag_name, order="count")
|
|
found = [t for t in tags if t["name"] == tag_name]
|
|
assert found, f"Tag '{tag_name}' not found in API response"
|
|
t = found[0]
|
|
cat = CATEGORY_NAMES.get(t.get("category", -1), "unknown")
|
|
print(f" {tag_name:20s} category={cat:12s} posts={t['post_count']:>8,}")
|
|
print(" PASS ✓")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main() -> None:
|
|
env = load_env()
|
|
api_key = env.get("DANBOORU_API_KEY") or os.environ.get("DANBOORU_API_KEY")
|
|
username = env.get("DANBOORU_USERNAME") or os.environ.get("DANBOORU_USERNAME")
|
|
|
|
if api_key:
|
|
print(f"API key loaded: {api_key[:8]}…")
|
|
else:
|
|
print("No API key found — using unauthenticated access (public endpoints only)")
|
|
|
|
# Danbooru public tag endpoints don't require auth.
|
|
# Pass username + api_key only when both are available.
|
|
session = make_session(
|
|
api_key=api_key if (api_key and username) else None,
|
|
username=username,
|
|
)
|
|
|
|
tests = [
|
|
test_basic_fetch,
|
|
test_fields_present,
|
|
test_pagination,
|
|
test_category_filter,
|
|
test_name_search,
|
|
test_well_known_tags,
|
|
]
|
|
|
|
passed = 0
|
|
failed = 0
|
|
for test_fn in tests:
|
|
try:
|
|
test_fn(session)
|
|
passed += 1
|
|
except Exception as exc:
|
|
print(f" FAIL ✗ {exc}")
|
|
failed += 1
|
|
time.sleep(0.3) # be polite to the API
|
|
|
|
print(f"\n{'='*50}")
|
|
print(f"Results: {passed} passed, {failed} failed")
|
|
if failed:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|