MediaShelf/tools/mockserver.py

#!/usr/bin/env python3
"""A fake Plex + Tautulli pair, good enough to exercise the whole ingest.

Used two ways:
  * standalone, to try tools/probe.py without touching a live server
  * imported by the test suite (tests/conftest.py) as a fixture

It deliberately reproduces the awkward parts of the real APIs: multi-version
movies, split parts, shows whose episodes carry no show GUID, Tautulli's
200-with-result-error failure mode, and the fact that a library's item count and
its reported totalSize have to agree for paging to terminate.

Not part of the application.
"""

from __future__ import annotations

import json
import random
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

NOW = int(time.time())
DAY = 86400

MACHINE_ID = "abc123def456abc123def456"
COVERAGE_DAYS = 400          # Tautulli has been logging this long


def _build(seed=7):
    rnd = random.Random(seed)
    movies, episodes, shows, history = [], [], [], []

    # ── movies, in two libraries that share GUIDs (the duplicate case) ──
    for i in range(60):
        guid = "plex://movie/%04d" % i
        parts = [{
            "id": 9000 + i, "file": "/mnt/titan4/Movies/Film %d/Film %d.mkv" % (i, i),
            "size": rnd.randint(2, 40) * 10**9, "container": "mkv",
        }]
        if i % 12 == 0:                       # a split file
            parts.append({
                "id": 9500 + i,
                "file": "/mnt/titan4/Movies/Film %d/Film %d.part2.mkv" % (i, i),
                "size": rnd.randint(1, 8) * 10**9, "container": "mkv",
            })
        movies.append({
            "sectionKey": "1", "ratingKey": str(1000 + i), "guid": guid, "type": "movie",
            "title": "Film %d" % i, "titleSort": "Film %04d" % i, "year": 1990 + i % 35,
            "addedAt": NOW - rnd.randint(10, 3000) * DAY,
            "updatedAt": NOW - rnd.randint(1, 500) * DAY,
            "duration": rnd.randint(80, 190) * 60000,
            "viewCount": rnd.choice([0, 0, 0, 1, 2]),
            "Media": [{"videoResolution": "1080", "videoCodec": "h264",
                       "audioCodec": "eac3", "bitrate": 8000, "Part": parts}],
        })

    # 4K copies of the first 15 films: same GUID, different library
    for i in range(15):
        movies.append({
            "sectionKey": "2", "ratingKey": str(2000 + i), "guid": "plex://movie/%04d" % i,
            "type": "movie", "title": "Film %d" % i, "year": 1990 + i % 35,
            "addedAt": NOW - rnd.randint(10, 900) * DAY,
            "duration": rnd.randint(80, 190) * 60000, "viewCount": rnd.choice([0, 1]),
            "Media": [{"videoResolution": "4k", "videoCodec": "hevc", "bitrate": 40000,
                       "Part": [{"id": 7000 + i,
                                 "file": "/mnt/titan4/4K Movies/Film %d/Film %d.4k.mkv" % (i, i),
                                 "size": rnd.randint(40, 90) * 10**9, "container": "mkv"}]}],
        })

    # Plex's addedAt follows the FILE: re-encode one and Date Added resets while
    # the item and its watch history survive. This movie reproduces that — added
    # "yesterday" but watched two years ago. 19% of the live library is like this.
    movies.append({
        "sectionKey": "1", "ratingKey": "1999", "guid": "plex://movie/replaced",
        "type": "movie", "title": "Replaced File", "year": 2001,
        "addedAt": NOW - 1 * DAY, "updatedAt": NOW - 1 * DAY,
        "duration": 100 * 60000, "viewCount": 3,
        "Media": [{"videoResolution": "1080", "Part": [
            {"id": 6999, "file": "/mnt/titan4/Movies/Replaced File/Replaced.mkv",
             "size": 9 * 10**9}]}],
    })

    # The case that actually failed on the live library: Plex says it was added
    # last week, Plex ALSO says it was last watched years ago, and Tautulli has
    # no record at all because its history only starts recently. Only
    # lastViewedAt carries the evidence here.
    movies.append({
        "sectionKey": "1", "ratingKey": "1998", "guid": "plex://movie/plexonly",
        "type": "movie", "title": "Plex Evidence Only", "year": 1999,
        "addedAt": NOW - 7 * DAY, "updatedAt": NOW - 7 * DAY,
        "lastViewedAt": NOW - 1500 * DAY,        # long before Tautulli existed
        "duration": 100 * 60000, "viewCount": 2,
        "Media": [{"videoResolution": "1080", "Part": [
            {"id": 6998, "file": "/mnt/titan4/Movies/Plex Evidence Only/x.mkv",
             "size": 7 * 10**9}]}],
    })

    # ── shows / seasons / episodes ─────────────────────────────────────
    ep_rk = 50000
    for s in range(6):
        show_rk = str(900 + s)
        shows.append({
            "sectionKey": "3", "ratingKey": show_rk, "guid": "plex://show/%04d" % s,
            "type": "show", "title": "Show %d" % s,
        })
        for se in range(1, rnd.randint(2, 4)):
            season_rk = str(9000 + s * 10 + se)
            for ep in range(1, rnd.randint(6, 13)):
                episodes.append({
                    "sectionKey": "3", "ratingKey": str(ep_rk), "guid": "plex://episode/%d" % ep_rk,
                    "type": "episode", "title": "Episode %d" % ep, "index": ep,
                    "grandparentRatingKey": show_rk, "grandparentTitle": "Show %d" % s,
                    "parentRatingKey": season_rk, "parentIndex": se,
                    "addedAt": NOW - rnd.randint(10, 2500) * DAY,
                    "duration": rnd.randint(20, 55) * 60000,
                    "viewCount": rnd.choice([0, 0, 1]),
                    "Media": [{"videoResolution": "1080", "videoCodec": "h264",
                               "Part": [{"id": 60000 + ep_rk,
                                         "file": "/mnt/titan4/TVShows/Show %d/S%02d/E%02d.mkv" % (s, se, ep),
                                         "size": rnd.randint(1, 5) * 10**9}]}],
                })
                ep_rk += 1

    # A show whose episodes carry no parentRatingKey — Plex really does this
    # (Firefly, in the live TV Show Archive). The season must still be built.
    orphan_show = "990"
    shows.append({"sectionKey": "3", "ratingKey": orphan_show,
                  "guid": "plex://show/orphan", "type": "show", "title": "Orphan Show"})
    for ep in range(1, 6):
        episodes.append({
            "sectionKey": "3", "ratingKey": str(ep_rk), "guid": "plex://episode/%d" % ep_rk,
            "type": "episode", "title": "Orphan Ep %d" % ep, "index": ep,
            "grandparentRatingKey": orphan_show, "grandparentTitle": "Orphan Show",
            "parentRatingKey": None, "parentIndex": 1,
            "addedAt": NOW - 900 * DAY, "duration": 45 * 60000, "viewCount": 0,
            "Media": [{"videoResolution": "1080", "Part": [
                {"id": 70000 + ep_rk,
                 "file": "/mnt/titan4/TVArchive/Orphan Show/S01/E%02d.mkv" % ep,
                 "size": 2 * 10**9}]}],
        })
        ep_rk += 1

    # ── Tautulli history ───────────────────────────────────────────────
    # 1998/1999 are the date-correction fixtures: their history is set
    # explicitly below, so keep the random plays away from them.
    SPECIAL = {"1998", "1999"}
    watchable = [m["ratingKey"] for m in movies if m["ratingKey"] not in SPECIAL] \
                + [e["ratingKey"] for e in episodes]
    row_id = 1
    for _ in range(1200):
        rk = rnd.choice(watchable)
        pc = rnd.choice([2, 6, 11, 30, 55, 70, 88, 95, 99, 100, 100])
        started = NOW - rnd.randint(1, COVERAGE_DAYS) * DAY - rnd.randint(0, 80000)
        history.append({
            "row_id": row_id, "reference_id": row_id // 2,
            "date": started, "started": started,
            "stopped": started + rnd.randint(300, 7200),
            "rating_key": rk,
            "user_id": rnd.choice([1, 2, 3, 4, 5]),
            "user": "user%d" % rnd.choice([1, 2, 3]),
            "friendly_name": rnd.choice(["Jess", "Sam", "Alex", "Guest"]),
            "media_type": "movie" if rk.startswith(("1", "2")) else "episode",
            "percent_complete": pc,
            "watched_status": 1 if pc >= 85 else (0.5 if pc >= 15 else 0),
            "play_duration": rnd.randint(120, 7000), "paused_counter": 0,
            "platform": "Chrome",
        })
        row_id += 1
    # guarantee the oldest event is exactly COVERAGE_DAYS old, so pre_history is testable
    for k in range(3):
        t = NOW - (300 - k * 20) * DAY
        history.append({
            "row_id": row_id + k, "reference_id": (row_id + k) // 2,
            "date": t, "started": t, "stopped": t + 6000,
            "rating_key": "1999", "user_id": 1, "user": "user1",
            "friendly_name": "Jess", "media_type": "movie",
            "percent_complete": 99, "watched_status": 1,
            "play_duration": 6000, "paused_counter": 0, "platform": "Chrome",
        })
    history[0]["date"] = history[0]["started"] = NOW - COVERAGE_DAYS * DAY
    history.sort(key=lambda h: h["date"])
    return movies, episodes, shows, history


MOVIES, EPISODES, SHOWS, HISTORY = _build()

SECTIONS = [
    {"key": "1", "type": "movie", "title": "Movies",
     "Location": [{"path": "/mnt/titan4/Movies"}]},
    {"key": "2", "type": "movie", "title": "4K Movies",
     "Location": [{"path": "/mnt/titan4/4K Movies"}]},
    {"key": "3", "type": "show", "title": "TV Shows",
     "Location": [{"path": "/mnt/titan4/TVShows"}]},
    {"key": "4", "type": "artist", "title": "Music",
     "Location": [{"path": "/mnt/titan4/Music"}]},
]


def _page(items, headers, query):
    start = int(headers.get("X-Plex-Container-Start")
                or query.get("X-Plex-Container-Start", ["0"])[0])
    size = int(headers.get("X-Plex-Container-Size")
               or query.get("X-Plex-Container-Size", [str(len(items))])[0])
    return items[start:start + size], len(items)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *a):
        pass

    def _send(self, obj, code=200):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):  # noqa: N802
        u = urllib.parse.urlsplit(self.path)
        q = urllib.parse.parse_qs(u.query)

        # ── Plex ───────────────────────────────────────────────────────
        if u.path == "/identity":
            return self._send({"MediaContainer": {
                "machineIdentifier": MACHINE_ID, "version": "1.43.3.10861",
                "friendlyName": "Loki"}})

        if u.path == "/library/sections":
            return self._send({"MediaContainer": {"Directory": SECTIONS}})

        if u.path == "/accounts":
            return self._send({"MediaContainer": {"Account": [
                {"id": 1, "name": "jess"}, {"id": 2, "name": "sam"}]}})

        if u.path.startswith("/library/sections/") and u.path.endswith("/all"):
            key = u.path.split("/")[3]
            libtype = q.get("type", ["1"])[0]
            if libtype == "1":
                src = [m for m in MOVIES if m["sectionKey"] == key]
            elif libtype == "4":
                src = [e for e in EPISODES if e["sectionKey"] == key]
            elif libtype == "2":
                src = [s for s in SHOWS if s["sectionKey"] == key]
            else:
                src = []
            batch, total = _page(src, self.headers, q)
            return self._send({"MediaContainer": {
                "Metadata": batch, "totalSize": total, "size": len(batch)}})

        if u.path == "/status/sessions/history/all":
            rows = [{"historyKey": "/status/sessions/history/%d" % h["row_id"],
                     "ratingKey": h["rating_key"], "viewedAt": h["date"],
                     "accountID": h["user_id"], "type": h["media_type"]}
                    for h in sorted(HISTORY, key=lambda x: -x["date"])]
            batch, total = _page(rows, self.headers, q)
            return self._send({"MediaContainer": {
                "Metadata": batch, "totalSize": total, "size": len(batch)}})

        # ── Tautulli ───────────────────────────────────────────────────
        if u.path == "/api/v2":
            cmd = q.get("cmd", [""])[0]
            if q.get("apikey", [""])[0] == "BADKEY":
                return self._send({"response": {"result": "error",
                                                "message": "Invalid apikey", "data": None}})
            if cmd == "get_server_info":
                ident = q.get("_force_id", [MACHINE_ID])[0]
                return self._send({"response": {"result": "success", "data": {
                    "pms_identifier": ident, "pms_name": "Loki",
                    "pms_ip": "192.168.1.10", "pms_port": 32400}}})
            if cmd == "get_users":
                return self._send({"response": {"result": "success", "data": [
                    {"user_id": i, "username": "user%d" % i,
                     "friendly_name": n} for i, n in
                    enumerate(["Jess", "Sam", "Alex", "Guest", "Kid"], start=1)]}})
            if cmd == "get_history":
                direction = q.get("order_dir", ["desc"])[0]
                rows = sorted(HISTORY, key=lambda h: h["date"],
                              reverse=(direction != "asc"))
                start = int(q.get("start", ["0"])[0])
                length = int(q.get("length", ["25"])[0])
                return self._send({"response": {"result": "success", "data": {
                    "data": rows[start:start + length],
                    "recordsFiltered": len(rows), "recordsTotal": len(rows)}}})
            if cmd == "get_libraries":
                return self._send({"response": {"result": "success", "data": [
                    {"section_id": s["key"], "section_name": s["title"],
                     "section_type": s["type"]} for s in SECTIONS]}})
            if cmd == "get_library_media_info":
                # Show sections really do return zero sizes here (§4.11).
                key = q.get("section_id", [""])[0]
                sec = next((s for s in SECTIONS if s["key"] == key), None)
                if sec and sec["type"] == "show":
                    rows = [{"rating_key": s["ratingKey"], "title": s["title"],
                             "media_type": "show", "file_size": 0, "play_count": 0}
                            for s in SHOWS]
                else:
                    rows = [{"rating_key": m["ratingKey"], "title": m["title"],
                             "media_type": "movie",
                             "file_size": sum(p["size"] for md in m["Media"] for p in md["Part"]),
                             "added_at": m["addedAt"], "play_count": m.get("viewCount", 0)}
                            for m in MOVIES if m["sectionKey"] == key]
                return self._send({"response": {"result": "success", "data": {
                    "data": rows, "recordsFiltered": len(rows), "recordsTotal": len(rows)}}})
            return self._send({"response": {"result": "error",
                                            "message": "unknown cmd", "data": None}})

        self.send_response(404)
        self.send_header("Content-Length", "0")
        self.end_headers()


class FakeServer:
    """Context manager wrapping the handler on an ephemeral port."""

    def __init__(self, port: int = 0):
        self.httpd = ThreadingHTTPServer(("127.0.0.1", port), Handler)
        self.httpd.daemon_threads = True
        self.port = self.httpd.server_address[1]
        self.url = "http://127.0.0.1:%d" % self.port
        self._thread: threading.Thread | None = None

    def __enter__(self) -> "FakeServer":
        self._thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, *exc):
        self.httpd.shutdown()
        self.httpd.server_close()
        return False


def stats() -> dict:
    return {
        "movies": len(MOVIES), "episodes": len(EPISODES), "shows": len(SHOWS),
        "history": len(HISTORY),
        "total_bytes": sum(p["size"] for m in MOVIES for md in m["Media"] for p in md["Part"])
                       + sum(p["size"] for e in EPISODES for md in e["Media"] for p in md["Part"]),
    }


if __name__ == "__main__":
    srv = FakeServer(8899)
    print("mock Plex+Tautulli on %s" % srv.url)
    print(json.dumps(stats(), indent=2))
    srv.httpd.serve_forever()