Initial commit
This commit is contained in:
Aodhan Collins
221
scripts/test_danbooru_api.py
Normal file
221
scripts/test_danbooru_api.py
Normal file
@@ -0,0 +1,221 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user