Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6d7f26f7be
|
||
|
|
446c0cd9ca
|
||
|
|
7b395e1834
|
||
|
|
b1625bfab5
|
||
|
|
b22b1ffab5
|
||
|
|
de19074fc0
|
||
|
|
ef0d1ce144
|
||
|
|
85c71faf2d
|
||
|
|
362541e13b
|
||
|
|
920f9f9f2d
|
||
|
|
23172111e9
|
||
|
|
b680f52126
|
||
|
|
eb81fa478d
|
||
|
|
68975fcadd
|
||
|
|
0b6e292cbd
|
||
|
|
68e0bc190c
|
||
|
|
c577dc5267
|
||
|
|
abe1444e4c
|
||
|
|
603a927f12
|
||
|
|
2dc8206f6f
|
||
|
|
2b65dfe076
|
||
|
|
0013726b7f
|
||
|
|
b8f509ac63
|
||
|
|
b0fdf12808
|
||
|
|
cf46b71883
|
||
|
|
bc9f7fa539
|
||
|
|
70095f3184
|
||
|
|
0fa51d6199
|
||
|
|
ca3011011a
|
||
|
|
53b8bd3fae
|
||
|
|
093426ac61
|
||
|
|
c1f967ece9
|
||
|
|
23cc721612
|
||
|
|
a23d252f27
|
||
|
|
2cf00fde4f
|
||
|
|
371f7d2ca3
|
||
|
|
12e5e92fe4
|
||
|
|
a7b4337835
|
||
|
|
a8d3a5c421
|
||
|
|
6aa8527adc
|
||
|
|
56d3c36c7d
|
||
|
|
2d13dfb063
|
||
|
|
34ae3b3e64
|
||
|
|
51f33a8dac
|
||
|
|
ee73e36f8e
|
||
|
|
0f2f43a55a
|
||
|
|
fb92bc6766
|
@@ -0,0 +1,35 @@
|
|||||||
|
name: Audit
|
||||||
|
|
||||||
|
on:
|
||||||
|
schedule:
|
||||||
|
- cron: "0 0 * * 1"
|
||||||
|
push:
|
||||||
|
paths: [uv.lock]
|
||||||
|
pull_request:
|
||||||
|
paths: [uv.lock]
|
||||||
|
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
|
||||||
|
env:
|
||||||
|
UV_PYTHON_DOWNLOADS: never
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
audit:
|
||||||
|
name: Dependencies
|
||||||
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
|
steps:
|
||||||
|
- name: Checkout repository
|
||||||
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
|
|
||||||
|
- name: Cache uv packages
|
||||||
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
|
with:
|
||||||
|
path: ~/.cache/uv
|
||||||
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
|
- name: Install dependencies
|
||||||
|
run: uv sync --locked
|
||||||
|
|
||||||
|
- name: Audit dependencies with pip-audit
|
||||||
|
run: uv run pip-audit --skip-editable
|
||||||
+39
-8
@@ -4,29 +4,60 @@ on:
|
|||||||
push:
|
push:
|
||||||
tags: ["v*"]
|
tags: ["v*"]
|
||||||
|
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
# packages: write # not yet supported by Gitea
|
||||||
|
|
||||||
|
env:
|
||||||
|
UV_PYTHON_DOWNLOADS: never
|
||||||
|
|
||||||
jobs:
|
jobs:
|
||||||
publish:
|
publish:
|
||||||
name: Publish
|
name: Publish
|
||||||
runs-on: logaldeveloper-archlinux
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
with:
|
with:
|
||||||
fetch-depth: 0
|
fetch-depth: 0
|
||||||
|
|
||||||
- name: Cache uv packages
|
|
||||||
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
|
|
||||||
with:
|
|
||||||
path: ~/.cache/uv
|
|
||||||
key: uv-${{ hashFiles('uv.lock') }}
|
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: uv sync
|
run: uv sync --locked --no-dev
|
||||||
|
|
||||||
|
- name: Generate package metadata
|
||||||
|
id: metadata
|
||||||
|
run: |
|
||||||
|
version="${{ gitea.ref_name }}"
|
||||||
|
version="${version#v}"
|
||||||
|
printf 'version=%s\n' "$version" | tee -a "$GITEA_OUTPUT"
|
||||||
|
|
||||||
- name: Build package
|
- name: Build package
|
||||||
run: uv build
|
run: uv build
|
||||||
|
|
||||||
|
- name: Generate SBOM
|
||||||
|
env:
|
||||||
|
SYFT_CHECK_FOR_APP_UPDATE: "false"
|
||||||
|
run: |
|
||||||
|
syft scan dir:.venv \
|
||||||
|
--override-default-catalogers python-installed-package-cataloger \
|
||||||
|
--select-catalogers=-file \
|
||||||
|
--source-name git.logal.dev/LogalDeveloper/Crabstero \
|
||||||
|
--source-version "${{ steps.metadata.outputs.version }}" \
|
||||||
|
--output syft-table \
|
||||||
|
--output cyclonedx-json=crabstero-${{ steps.metadata.outputs.version }}.cyclonedx.json
|
||||||
|
sha256sum crabstero-${{ steps.metadata.outputs.version }}.cyclonedx.json
|
||||||
|
zstd -T0 --ultra -22 \
|
||||||
|
crabstero-${{ steps.metadata.outputs.version }}.cyclonedx.json
|
||||||
|
|
||||||
- name: Publish package
|
- name: Publish package
|
||||||
run: uv publish
|
run: uv publish
|
||||||
env:
|
env:
|
||||||
UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
|
UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
|
||||||
|
|
||||||
|
- name: Upload SBOM artifact
|
||||||
|
uses: christopherhx/gitea-upload-artifact@8818363695ca2d5782c64f6453273341374767b7 # v7
|
||||||
|
with:
|
||||||
|
name: crabstero-cyclonedx-${{ steps.metadata.outputs.version }}
|
||||||
|
path: crabstero-${{ steps.metadata.outputs.version }}.cyclonedx.json.zst
|
||||||
|
if-no-files-found: error
|
||||||
|
archive: "false"
|
||||||
|
|||||||
+37
-21
@@ -2,69 +2,68 @@ name: CI
|
|||||||
|
|
||||||
on:
|
on:
|
||||||
push:
|
push:
|
||||||
branches: [master]
|
|
||||||
pull_request:
|
pull_request:
|
||||||
|
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
|
||||||
|
env:
|
||||||
|
UV_PYTHON_DOWNLOADS: never
|
||||||
|
|
||||||
jobs:
|
jobs:
|
||||||
formatting:
|
formatting:
|
||||||
name: Formatting
|
name: Formatting
|
||||||
runs-on: logaldeveloper-archlinux
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
with:
|
|
||||||
fetch-depth: 0
|
|
||||||
|
|
||||||
- name: Cache uv packages
|
- name: Cache uv packages
|
||||||
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
with:
|
with:
|
||||||
path: ~/.cache/uv
|
path: ~/.cache/uv
|
||||||
key: uv-${{ hashFiles('uv.lock') }}
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: uv sync
|
run: uv sync --locked
|
||||||
|
|
||||||
- name: Check formatting with Ruff
|
- name: Check formatting with Ruff
|
||||||
run: uv run ruff format --check --diff .
|
run: uv run ruff format --check --diff .
|
||||||
|
|
||||||
linting:
|
linting:
|
||||||
name: Linting
|
name: Linting
|
||||||
runs-on: logaldeveloper-archlinux
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
with:
|
|
||||||
fetch-depth: 0
|
|
||||||
|
|
||||||
- name: Cache uv packages
|
- name: Cache uv packages
|
||||||
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
with:
|
with:
|
||||||
path: ~/.cache/uv
|
path: ~/.cache/uv
|
||||||
key: uv-${{ hashFiles('uv.lock') }}
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: uv sync
|
run: uv sync --locked
|
||||||
|
|
||||||
- name: Check linting with Ruff
|
- name: Check linting with Ruff
|
||||||
run: uv run ruff check .
|
run: uv run ruff check .
|
||||||
|
|
||||||
tests:
|
tests:
|
||||||
name: Tests
|
name: Tests
|
||||||
runs-on: logaldeveloper-archlinux
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
with:
|
|
||||||
fetch-depth: 0
|
|
||||||
|
|
||||||
- name: Cache uv packages
|
- name: Cache uv packages
|
||||||
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
with:
|
with:
|
||||||
path: ~/.cache/uv
|
path: ~/.cache/uv
|
||||||
key: uv-${{ hashFiles('uv.lock') }}
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: uv sync
|
run: uv sync --locked
|
||||||
|
|
||||||
- name: Run unit tests with Pytest
|
- name: Run unit tests with Pytest
|
||||||
run: uv run pytest -v --cov --cov-report=
|
run: uv run pytest -v --cov --cov-report=
|
||||||
@@ -74,21 +73,38 @@ jobs:
|
|||||||
|
|
||||||
type-checking:
|
type-checking:
|
||||||
name: Type Checking
|
name: Type Checking
|
||||||
runs-on: logaldeveloper-archlinux
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
with:
|
|
||||||
fetch-depth: 0
|
|
||||||
|
|
||||||
- name: Cache uv packages
|
- name: Cache uv packages
|
||||||
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
with:
|
with:
|
||||||
path: ~/.cache/uv
|
path: ~/.cache/uv
|
||||||
key: uv-${{ hashFiles('uv.lock') }}
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: uv sync
|
run: uv sync --locked
|
||||||
|
|
||||||
- name: Check types with Mypy
|
- name: Check types with Mypy
|
||||||
run: uv run mypy .
|
run: uv run mypy .
|
||||||
|
|
||||||
|
spelling:
|
||||||
|
name: Spelling
|
||||||
|
runs-on: logaldeveloper-archlinux-ci
|
||||||
|
steps:
|
||||||
|
- name: Checkout repository
|
||||||
|
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||||
|
|
||||||
|
- name: Cache uv packages
|
||||||
|
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
|
||||||
|
with:
|
||||||
|
path: ~/.cache/uv
|
||||||
|
key: uv-${{ hashFiles('uv.lock') }}
|
||||||
|
|
||||||
|
- name: Install dependencies
|
||||||
|
run: uv sync --locked
|
||||||
|
|
||||||
|
- name: Check spelling with codespell
|
||||||
|
run: uv run codespell
|
||||||
|
|||||||
@@ -11,6 +11,8 @@ dist/
|
|||||||
build/
|
build/
|
||||||
*.egg-info/
|
*.egg-info/
|
||||||
*.egg
|
*.egg
|
||||||
|
/*.cyclonedx.json
|
||||||
|
/*.cyclonedx.json.zst
|
||||||
|
|
||||||
# Virtual environments
|
# Virtual environments
|
||||||
venv/
|
venv/
|
||||||
@@ -23,6 +25,9 @@ env/
|
|||||||
*.swp
|
*.swp
|
||||||
*.swo
|
*.swo
|
||||||
|
|
||||||
|
# Testing
|
||||||
|
.coverage
|
||||||
|
|
||||||
# Runtime data
|
# Runtime data
|
||||||
*.db
|
*.db
|
||||||
.env
|
.env
|
||||||
|
|||||||
+179
-24
@@ -12,28 +12,79 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Entry point for the Crabstero Discord bot.
|
||||||
Entry point for the Crabstero Discord bot.
|
|
||||||
|
|
||||||
Provides an argparse CLI with environment variable and systemd credential fallbacks for --token and --database-path.
|
Provides an argparse CLI with environment variable and systemd credential
|
||||||
|
fallbacks, plus systemd notify/watchdog integration when run as a service.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import contextlib
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import signal
|
import signal
|
||||||
|
import socket
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import uvloop
|
||||||
|
|
||||||
from crabstero import __version__ as crabstero_version
|
from crabstero import __version__ as crabstero_version
|
||||||
from crabstero.bot import Crabstero
|
from crabstero.bot import Crabstero
|
||||||
|
from crabstero.metrics import TcpMetricsAddress, UnixMetricsAddress
|
||||||
|
|
||||||
logger = logging.getLogger("crabstero")
|
logger = logging.getLogger("crabstero")
|
||||||
|
|
||||||
|
|
||||||
def _read_credential(name: str) -> str | None:
|
def _sd_notify(state: str) -> None:
|
||||||
|
"""Send a notification to systemd, if running under a systemd service.
|
||||||
|
|
||||||
|
:param state: The systemd notification payload.
|
||||||
"""
|
"""
|
||||||
Read a value from a systemd credential file.
|
addr = os.environ.get("NOTIFY_SOCKET")
|
||||||
|
if not addr:
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
|
||||||
|
if addr[0] == "@":
|
||||||
|
addr = "\0" + addr[1:]
|
||||||
|
sock.sendto(state.encode(), addr)
|
||||||
|
except OSError as exc:
|
||||||
|
logger.debug("Could not send systemd notification: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
|
def _watchdog_interval() -> float | None:
|
||||||
|
"""Return the systemd watchdog ping interval in seconds, if enabled."""
|
||||||
|
usec = os.environ.get("WATCHDOG_USEC")
|
||||||
|
if not usec:
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
watchdog_usec = int(usec)
|
||||||
|
except ValueError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if watchdog_usec <= 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return watchdog_usec / 1_000_000 / 2
|
||||||
|
|
||||||
|
|
||||||
|
async def _watchdog_loop(interval: float, stop: asyncio.Event) -> None:
|
||||||
|
"""Send systemd watchdog pings until *stop* is set.
|
||||||
|
|
||||||
|
:param interval: Seconds to wait between watchdog pings.
|
||||||
|
:param stop: Event that stops the watchdog loop.
|
||||||
|
"""
|
||||||
|
while not stop.is_set():
|
||||||
|
_sd_notify("WATCHDOG=1")
|
||||||
|
with contextlib.suppress(TimeoutError):
|
||||||
|
await asyncio.wait_for(stop.wait(), timeout=interval)
|
||||||
|
|
||||||
|
|
||||||
|
def _read_credential(name: str) -> str | None:
|
||||||
|
"""Read a value from a systemd credential file.
|
||||||
|
|
||||||
Looks for a file named *name* inside the directory pointed to by the
|
Looks for a file named *name* inside the directory pointed to by the
|
||||||
``CREDENTIALS_DIRECTORY`` environment variable (set automatically by
|
``CREDENTIALS_DIRECTORY`` environment variable (set automatically by
|
||||||
@@ -53,8 +104,7 @@ def _read_credential(name: str) -> str | None:
|
|||||||
|
|
||||||
|
|
||||||
def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||||
"""
|
"""Parse command-line arguments with environment variable fallbacks.
|
||||||
Parses command-line arguments with environment variable fallbacks.
|
|
||||||
|
|
||||||
:param argv: Optional argument list (defaults to sys.argv[1:]).
|
:param argv: Optional argument list (defaults to sys.argv[1:]).
|
||||||
:return: Parsed arguments namespace.
|
:return: Parsed arguments namespace.
|
||||||
@@ -66,35 +116,77 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--token",
|
"--token",
|
||||||
default=os.environ.get("TOKEN") or _read_credential("token"),
|
default=os.environ.get("TOKEN") or _read_credential("token"),
|
||||||
help="Discord bot token (default: TOKEN environment variable or systemd credential 'token').",
|
help=(
|
||||||
|
"Discord bot token"
|
||||||
|
" (default: TOKEN environment variable"
|
||||||
|
" or systemd credential 'token')."
|
||||||
|
),
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--database-path",
|
"--database-path",
|
||||||
"--database",
|
"--database",
|
||||||
default=os.environ.get("DATABASE_PATH", "crabstero.db"),
|
default=os.environ.get("DATABASE_PATH", "crabstero.db"),
|
||||||
help='Path to the SQLite database file (default: DATABASE_PATH environment variable or "crabstero.db").',
|
help=(
|
||||||
|
"Path to the SQLite database file"
|
||||||
|
" (default: DATABASE_PATH environment"
|
||||||
|
' variable or "crabstero.db").'
|
||||||
|
),
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--ingestion-workers",
|
"--ingest-only",
|
||||||
type=int,
|
action="store_true",
|
||||||
default=int(os.environ.get("INGESTION_WORKERS", "4")),
|
default=False,
|
||||||
help="Number of concurrent ingestion workers (default: INGESTION_WORKERS environment variable or 4).",
|
help=(
|
||||||
|
"Run in ingest-only mode: ingest channel"
|
||||||
|
" history and real-time messages but"
|
||||||
|
" never respond."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--listen-metrics",
|
||||||
|
default=os.environ.get("LISTEN_METRICS"),
|
||||||
|
help=(
|
||||||
|
"Enable Prometheus metrics endpoint on HOST:PORT or unix:/path.sock"
|
||||||
|
" (e.g. 127.0.0.1:9090 or unix:/run/crabstero.sock)."
|
||||||
|
" Disabled by default."
|
||||||
|
" (default: LISTEN_METRICS environment variable)."
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
args = parser.parse_args(argv)
|
args = parser.parse_args(argv)
|
||||||
|
|
||||||
|
if args.listen_metrics is not None:
|
||||||
|
if args.listen_metrics.startswith("unix:"):
|
||||||
|
path = args.listen_metrics.removeprefix("unix:")
|
||||||
|
if not path:
|
||||||
|
parser.error("--listen-metrics Unix socket path must not be empty")
|
||||||
|
args.listen_metrics = UnixMetricsAddress(path)
|
||||||
|
else:
|
||||||
|
host, sep, port_str = args.listen_metrics.rpartition(":")
|
||||||
|
if not sep or not host:
|
||||||
|
parser.error(
|
||||||
|
"--listen-metrics must be in HOST:PORT or unix:/path.sock format "
|
||||||
|
"(e.g. 127.0.0.1:9090)",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
args.listen_metrics = TcpMetricsAddress(host, int(port_str))
|
||||||
|
except ValueError:
|
||||||
|
parser.error(
|
||||||
|
f"--listen-metrics port must be an integer, got '{port_str}'",
|
||||||
|
)
|
||||||
|
|
||||||
if args.token is None:
|
if args.token is None:
|
||||||
parser.error(
|
parser.error(
|
||||||
"a Discord bot token is required via --token, the TOKEN environment variable, or a systemd credential named 'token'"
|
"a Discord bot token is required via --token,"
|
||||||
|
" the TOKEN environment variable,"
|
||||||
|
" or a systemd credential named 'token'",
|
||||||
)
|
)
|
||||||
|
|
||||||
return args
|
return args
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
"""
|
"""Run the bot. Parse arguments, configure logging, and start the bot."""
|
||||||
Main entry point. Parses arguments, configures logging, and starts the bot.
|
|
||||||
"""
|
|
||||||
args = _parse_args()
|
args = _parse_args()
|
||||||
|
|
||||||
logging.basicConfig(
|
logging.basicConfig(
|
||||||
@@ -107,19 +199,82 @@ def main() -> None:
|
|||||||
bot = Crabstero(
|
bot = Crabstero(
|
||||||
token=args.token,
|
token=args.token,
|
||||||
database_path=args.database_path,
|
database_path=args.database_path,
|
||||||
ingestion_workers=args.ingestion_workers,
|
ingest_only=args.ingest_only,
|
||||||
|
metrics_address=args.listen_metrics,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _run() -> None:
|
async def _run() -> None:
|
||||||
async with bot:
|
loop = asyncio.get_running_loop()
|
||||||
await bot.start()
|
watchdog_stop = asyncio.Event()
|
||||||
|
watchdog_task: asyncio.Task[None] | None = None
|
||||||
|
shutdown_task: asyncio.Task[None] | None = None
|
||||||
|
ready_notified = False
|
||||||
|
stopping_notified = False
|
||||||
|
|
||||||
# Make SIGTERM behave like SIGINT so systemd stop triggers the same
|
def notify_stopping() -> None:
|
||||||
# clean shutdown path (context manager __aexit__ -> bot.close()).
|
nonlocal stopping_notified
|
||||||
signal.signal(signal.SIGTERM, signal.default_int_handler)
|
if stopping_notified:
|
||||||
|
return
|
||||||
|
stopping_notified = True
|
||||||
|
_sd_notify("STOPPING=1")
|
||||||
|
|
||||||
|
async def shutdown() -> None:
|
||||||
|
logger.info("Received shutdown signal.")
|
||||||
|
notify_stopping()
|
||||||
|
await bot.close()
|
||||||
|
|
||||||
|
def request_shutdown() -> None:
|
||||||
|
nonlocal shutdown_task
|
||||||
|
if shutdown_task is None or shutdown_task.done():
|
||||||
|
shutdown_task = asyncio.create_task(
|
||||||
|
shutdown(),
|
||||||
|
name="crabstero-shutdown",
|
||||||
|
)
|
||||||
|
|
||||||
|
handled_signals = (signal.SIGINT, signal.SIGTERM)
|
||||||
|
for sig in handled_signals:
|
||||||
|
loop.add_signal_handler(sig, request_shutdown)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
asyncio.run(_run())
|
async with bot:
|
||||||
|
try:
|
||||||
|
await bot.login(args.token)
|
||||||
|
if bot.is_closed() or stopping_notified:
|
||||||
|
return
|
||||||
|
|
||||||
|
_sd_notify("READY=1")
|
||||||
|
ready_notified = True
|
||||||
|
|
||||||
|
wd_interval = _watchdog_interval()
|
||||||
|
if wd_interval is not None:
|
||||||
|
logger.info(
|
||||||
|
"systemd watchdog enabled (pinging every %.1fs).",
|
||||||
|
wd_interval,
|
||||||
|
)
|
||||||
|
watchdog_task = asyncio.create_task(
|
||||||
|
_watchdog_loop(wd_interval, watchdog_stop),
|
||||||
|
name="systemd-watchdog",
|
||||||
|
)
|
||||||
|
|
||||||
|
await bot.connect(reconnect=True)
|
||||||
|
finally:
|
||||||
|
if ready_notified:
|
||||||
|
notify_stopping()
|
||||||
|
|
||||||
|
watchdog_stop.set()
|
||||||
|
if watchdog_task is not None:
|
||||||
|
watchdog_task.cancel()
|
||||||
|
with contextlib.suppress(asyncio.CancelledError):
|
||||||
|
await watchdog_task
|
||||||
|
|
||||||
|
if shutdown_task is not None:
|
||||||
|
await shutdown_task
|
||||||
|
finally:
|
||||||
|
for sig in handled_signals:
|
||||||
|
loop.remove_signal_handler(sig)
|
||||||
|
|
||||||
|
try:
|
||||||
|
uvloop.run(_run())
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
logger.info("Interrupted.")
|
logger.info("Interrupted.")
|
||||||
|
|
||||||
|
|||||||
+241
-59
@@ -12,44 +12,63 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""The simple nonversation Discord bot.
|
||||||
The simple nonversation Discord bot.
|
|
||||||
|
|
||||||
Provides the Crabstero subclass that owns the full bot lifecycle: database connection,
|
Provides the Crabstero subclass that owns the full bot lifecycle: database connection,
|
||||||
cog loading, ingestion worker pool, and graceful shutdown.
|
cog loading, background ingestion, and graceful shutdown.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
|
from typing import override
|
||||||
|
|
||||||
import discord
|
import discord
|
||||||
from discord import app_commands
|
from discord import app_commands
|
||||||
|
from discord.app_commands import AppCommandError, CommandInvokeError
|
||||||
from discord.ext import commands
|
from discord.ext import commands
|
||||||
|
|
||||||
from crabstero import __version__ as crabstero_version
|
from crabstero import __version__ as crabstero_version
|
||||||
|
from crabstero.cache import IngestCache
|
||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
from crabstero.metrics import (
|
||||||
|
DISCORD_EVENTS,
|
||||||
|
DISCORD_LATENCY,
|
||||||
|
ERRORS,
|
||||||
|
GUILD_COUNT,
|
||||||
|
INGESTION_ACTIVE,
|
||||||
|
MetricsAddress,
|
||||||
|
MetricsServer,
|
||||||
|
)
|
||||||
from crabstero.tasks.ingestion import ingest_channel
|
from crabstero.tasks.ingestion import ingest_channel
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
type IngestableChannel = discord.TextChannel | discord.VoiceChannel
|
||||||
|
|
||||||
|
_MAX_CONCURRENT_INGESTIONS = 4
|
||||||
|
|
||||||
|
|
||||||
class Crabstero(commands.Bot):
|
class Crabstero(commands.Bot):
|
||||||
"""
|
"""Central bot subclass that owns all lifecycle state.
|
||||||
Central bot subclass that owns all lifecycle state.
|
|
||||||
|
|
||||||
The database is opened in setup_hook and closed in close(). Background
|
The database is opened in setup_hook and closed in close(). Background
|
||||||
ingestion is handled by a bounded queue and a fixed worker pool.
|
ingestion is handled by semaphore-bounded dynamic tasks.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, token: str, database_path: str, ingestion_workers: int = 4
|
self,
|
||||||
|
token: str,
|
||||||
|
database_path: str,
|
||||||
|
*,
|
||||||
|
ingest_only: bool = False,
|
||||||
|
metrics_address: MetricsAddress | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Configure intents, store configuration, and prepare ingestion state.
|
||||||
Configures intents, stores configuration, and prepares ingestion queue state.
|
|
||||||
|
|
||||||
:param token: The Discord bot token.
|
:param token: The Discord bot token.
|
||||||
:param database_path: The file path to the SQLite database.
|
:param database_path: The file path to the SQLite database.
|
||||||
:param ingestion_workers: The number of concurrent ingestion workers.
|
:param ingest_only: When True, the bot only ingests data and never responds.
|
||||||
|
:param metrics_address: Optional address for the Prometheus metrics server.
|
||||||
"""
|
"""
|
||||||
intents = discord.Intents.default()
|
intents = discord.Intents.default()
|
||||||
intents.guilds = True
|
intents.guilds = True
|
||||||
@@ -64,27 +83,106 @@ class Crabstero(commands.Bot):
|
|||||||
|
|
||||||
self._token = token
|
self._token = token
|
||||||
self._database_path = database_path
|
self._database_path = database_path
|
||||||
self._ingestion_worker_count = ingestion_workers
|
self._ingest_only = ingest_only
|
||||||
self._ingestion_queue: asyncio.Queue[
|
self._ingestion_semaphore = asyncio.Semaphore(_MAX_CONCURRENT_INGESTIONS)
|
||||||
discord.TextChannel | discord.VoiceChannel
|
self._ingestion_tasks: dict[int, asyncio.Task[None]] = {}
|
||||||
] = asyncio.Queue()
|
self._db: Database | None = None
|
||||||
self._ingestion_workers: list[asyncio.Task[None]] = []
|
self.ingest_cache = IngestCache()
|
||||||
self.db: Database
|
|
||||||
|
|
||||||
self.http.user_agent = f"DiscordBot (https://git.logal.dev/LogalDeveloper/Crabstero, {crabstero_version})"
|
self._metrics_address = metrics_address
|
||||||
|
self._metrics_server: MetricsServer | None = None
|
||||||
|
|
||||||
|
DISCORD_LATENCY.set_function(lambda: self.latency)
|
||||||
|
GUILD_COUNT.set_function(lambda: len(self.guilds))
|
||||||
|
INGESTION_ACTIVE.set_function(lambda: len(self._ingestion_tasks))
|
||||||
|
|
||||||
|
repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
|
||||||
|
self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def db(self) -> Database:
|
||||||
|
"""The active database connection.
|
||||||
|
|
||||||
|
:raises RuntimeError: If accessed before :meth:`setup_hook` has run.
|
||||||
|
"""
|
||||||
|
if self._db is None:
|
||||||
|
msg = "Database is not initialized"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
return self._db
|
||||||
|
|
||||||
|
@property
|
||||||
|
def ingest_only(self) -> bool:
|
||||||
|
"""Whether the bot is running in ingest-only mode."""
|
||||||
|
return self._ingest_only
|
||||||
|
|
||||||
|
@override
|
||||||
async def setup_hook(self) -> None:
|
async def setup_hook(self) -> None:
|
||||||
"""Opens the database, starts ingestion workers, loads all cogs, and syncs slash commands if changed."""
|
"""Open the database, start the ingest cache, and load all cogs."""
|
||||||
self.db = await Database.connect(self._database_path)
|
self._db = await Database.connect(self._database_path)
|
||||||
self._start_ingestion_workers()
|
self.ingest_cache.start()
|
||||||
|
|
||||||
from crabstero.listeners import interaction, message, server_events
|
if self._metrics_address is not None:
|
||||||
|
server = MetricsServer(self._metrics_address)
|
||||||
|
await server.start()
|
||||||
|
self._metrics_server = server
|
||||||
|
|
||||||
|
from crabstero.listeners import ( # noqa: PLC0415
|
||||||
|
interaction,
|
||||||
|
message,
|
||||||
|
server_events,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not self._ingest_only:
|
||||||
await interaction.setup(self)
|
await interaction.setup(self)
|
||||||
|
|
||||||
await message.setup(self)
|
await message.setup(self)
|
||||||
await server_events.setup(self)
|
await server_events.setup(self)
|
||||||
|
|
||||||
# Only sync slash commands if the registered commands differ from local definitions.
|
@self.tree.error
|
||||||
|
async def on_app_command_error(
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
error: AppCommandError,
|
||||||
|
) -> None:
|
||||||
|
original = (
|
||||||
|
error.original if isinstance(error, CommandInvokeError) else error
|
||||||
|
)
|
||||||
|
command_name = (
|
||||||
|
interaction.command.name if interaction.command else "unknown"
|
||||||
|
)
|
||||||
|
|
||||||
|
ERRORS.labels(source="command").inc()
|
||||||
|
logger.error(
|
||||||
|
"Unhandled exception in app command '%s'.",
|
||||||
|
command_name,
|
||||||
|
exc_info=original,
|
||||||
|
)
|
||||||
|
|
||||||
|
if self._metrics_server is not None:
|
||||||
|
reply = (
|
||||||
|
"I encountered an error while processing this command."
|
||||||
|
" The developer has been notified,"
|
||||||
|
" please try again later."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
reply = (
|
||||||
|
"I encountered an error while processing this command."
|
||||||
|
" Please try again later."
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
if interaction.response.is_done():
|
||||||
|
await interaction.followup.send(reply, ephemeral=True)
|
||||||
|
else:
|
||||||
|
await interaction.response.send_message(reply, ephemeral=True)
|
||||||
|
except discord.HTTPException:
|
||||||
|
logger.debug(
|
||||||
|
"Failed to send error response for command '%s'.",
|
||||||
|
command_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not self._ingest_only:
|
||||||
|
# Only sync slash commands if the registered commands
|
||||||
|
# differ from local definitions.
|
||||||
local_commands = {
|
local_commands = {
|
||||||
cmd.name: cmd.description
|
cmd.name: cmd.description
|
||||||
for cmd in self.tree.get_commands()
|
for cmd in self.tree.get_commands()
|
||||||
@@ -92,7 +190,8 @@ class Crabstero(commands.Bot):
|
|||||||
}
|
}
|
||||||
try:
|
try:
|
||||||
remote_commands = {
|
remote_commands = {
|
||||||
cmd.name: cmd.description for cmd in await self.tree.fetch_commands()
|
cmd.name: cmd.description
|
||||||
|
for cmd in await self.tree.fetch_commands()
|
||||||
}
|
}
|
||||||
except discord.HTTPException:
|
except discord.HTTPException:
|
||||||
remote_commands = {}
|
remote_commands = {}
|
||||||
@@ -101,53 +200,136 @@ class Crabstero(commands.Bot):
|
|||||||
logger.info("Slash command tree has changed, syncing with Discord.")
|
logger.info("Slash command tree has changed, syncing with Discord.")
|
||||||
await self.tree.sync()
|
await self.tree.sync()
|
||||||
|
|
||||||
|
@override
|
||||||
|
def dispatch(self, event: str, /, *args: object, **kwargs: object) -> None:
|
||||||
|
"""Dispatch an event, incrementing the events counter.
|
||||||
|
|
||||||
|
:param event: The event name.
|
||||||
|
"""
|
||||||
|
DISCORD_EVENTS.labels(event=event).inc()
|
||||||
|
super().dispatch(event, *args, **kwargs)
|
||||||
|
|
||||||
|
@override
|
||||||
|
async def on_error(
|
||||||
|
self,
|
||||||
|
event_method: str,
|
||||||
|
/,
|
||||||
|
*args: object,
|
||||||
|
**kwargs: object,
|
||||||
|
) -> None:
|
||||||
|
"""Increment the global error counter for event listener exceptions.
|
||||||
|
|
||||||
|
:param event_method: The name of the event that raised the exception.
|
||||||
|
"""
|
||||||
|
ERRORS.labels(source=event_method).inc()
|
||||||
|
logger.error("Unhandled exception in %s.", event_method, exc_info=True) # noqa: LOG014
|
||||||
|
|
||||||
async def on_ready(self) -> None:
|
async def on_ready(self) -> None:
|
||||||
"""Logs that the bot has started successfully."""
|
"""Log that the bot has started successfully."""
|
||||||
logger.info("Crabstero started!")
|
logger.info("Crabstero started!")
|
||||||
|
|
||||||
|
@override
|
||||||
async def start(self, token: str = "", *, reconnect: bool = True) -> None:
|
async def start(self, token: str = "", *, reconnect: bool = True) -> None:
|
||||||
"""
|
"""Start the bot using the token provided at initialization."""
|
||||||
Starts the bot using the stored token by default.
|
await super().start(self._token, reconnect=reconnect)
|
||||||
|
|
||||||
:param token: Optional token override. Falls back to the stored token if empty.
|
|
||||||
:param reconnect: Whether to automatically reconnect on disconnect.
|
|
||||||
"""
|
|
||||||
await super().start(token or self._token, reconnect=reconnect)
|
|
||||||
|
|
||||||
|
@override
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Cancels ingestion workers, closes the database, and then the bot connection."""
|
"""Cancel ingestion tasks, close the database, and then the bot connection."""
|
||||||
if self.is_closed():
|
if self.is_closed():
|
||||||
return
|
return
|
||||||
logger.info("Shutting down Crabstero...")
|
logger.info("Shutting down Crabstero...")
|
||||||
for worker in self._ingestion_workers:
|
tasks = list(self._ingestion_tasks.values())
|
||||||
worker.cancel()
|
self._ingestion_tasks.clear()
|
||||||
await asyncio.gather(*self._ingestion_workers, return_exceptions=True)
|
for task in tasks:
|
||||||
self._ingestion_workers.clear()
|
task.cancel()
|
||||||
if hasattr(self, "db"):
|
await asyncio.gather(*tasks, return_exceptions=True)
|
||||||
await self.db.close()
|
await self.ingest_cache.stop()
|
||||||
|
if self._metrics_server is not None:
|
||||||
|
await self._metrics_server.stop()
|
||||||
|
if self._db is not None:
|
||||||
|
await self._db.close()
|
||||||
await super().close()
|
await super().close()
|
||||||
|
|
||||||
def queue_channel_for_ingestion(
|
def queue_channel_for_ingestion(self, channel: IngestableChannel) -> None:
|
||||||
self, channel: discord.TextChannel | discord.VoiceChannel
|
"""Create a background task to ingest a single channel.
|
||||||
) -> None:
|
|
||||||
|
Duplicate requests for a channel that is already in-flight are ignored.
|
||||||
|
Concurrency is bounded by the ingestion semaphore.
|
||||||
|
|
||||||
|
:param channel: The channel to ingest.
|
||||||
"""
|
"""
|
||||||
Enqueues a single channel for background message history ingestion.
|
if channel.id in self._ingestion_tasks:
|
||||||
|
return
|
||||||
|
task = asyncio.create_task(
|
||||||
|
self._ingest_one(channel),
|
||||||
|
name=f"ingest-{channel.id}",
|
||||||
|
)
|
||||||
|
self._ingestion_tasks[channel.id] = task
|
||||||
|
task.add_done_callback(lambda t: self._on_ingestion_done(channel.id, t))
|
||||||
|
|
||||||
:param channel: The channel to enqueue.
|
async def _ingest_one(self, channel: IngestableChannel) -> None:
|
||||||
"""
|
"""Acquire the semaphore and ingest one channel."""
|
||||||
self._ingestion_queue.put_nowait(channel)
|
async with self._ingestion_semaphore:
|
||||||
|
|
||||||
def _start_ingestion_workers(self) -> None:
|
|
||||||
"""Spawns the fixed pool of ingestion worker tasks."""
|
|
||||||
for _ in range(self._ingestion_worker_count):
|
|
||||||
task = asyncio.create_task(self._ingestion_worker())
|
|
||||||
self._ingestion_workers.append(task)
|
|
||||||
|
|
||||||
async def _ingestion_worker(self) -> None:
|
|
||||||
"""Loops forever pulling channels from the ingestion queue and ingesting them."""
|
|
||||||
while True:
|
|
||||||
channel = await self._ingestion_queue.get()
|
|
||||||
try:
|
|
||||||
await ingest_channel(channel, self.db)
|
await ingest_channel(channel, self.db)
|
||||||
finally:
|
|
||||||
self._ingestion_queue.task_done()
|
def _on_ingestion_done(self, channel_id: int, task: asyncio.Task[None]) -> None:
|
||||||
|
"""Clean up a finished ingestion task and log any errors."""
|
||||||
|
self._ingestion_tasks.pop(channel_id, None)
|
||||||
|
if task.cancelled():
|
||||||
|
return
|
||||||
|
exc = task.exception()
|
||||||
|
if exc is not None:
|
||||||
|
ERRORS.labels(source="ingestion").inc()
|
||||||
|
logger.error("Ingestion task failed.", exc_info=exc)
|
||||||
|
|
||||||
|
|
||||||
|
class TrackedView(discord.ui.View):
|
||||||
|
"""Base View that increments the global error counter on failures."""
|
||||||
|
|
||||||
|
@override
|
||||||
|
async def on_error(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
error: Exception,
|
||||||
|
item: discord.ui.Item[TrackedView],
|
||||||
|
/,
|
||||||
|
) -> None:
|
||||||
|
"""Increment the error counter and log the exception.
|
||||||
|
|
||||||
|
:param interaction: The interaction that led to the failure.
|
||||||
|
:param error: The exception that was raised.
|
||||||
|
:param item: The item that failed the dispatch.
|
||||||
|
"""
|
||||||
|
ERRORS.labels(source="view").inc()
|
||||||
|
logger.error(
|
||||||
|
"Unhandled exception in view %r for item %r.",
|
||||||
|
self,
|
||||||
|
item,
|
||||||
|
exc_info=error,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TrackedModal(discord.ui.Modal):
|
||||||
|
"""Base Modal that increments the global error counter on failures."""
|
||||||
|
|
||||||
|
@override
|
||||||
|
async def on_error(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
error: Exception,
|
||||||
|
item: discord.ui.Item[TrackedModal] | None = None,
|
||||||
|
/,
|
||||||
|
) -> None:
|
||||||
|
"""Increment the error counter and log the exception.
|
||||||
|
|
||||||
|
:param interaction: The interaction that led to the failure.
|
||||||
|
:param error: The exception that was raised.
|
||||||
|
:param item: Unused. Present for BaseView signature compatibility.
|
||||||
|
"""
|
||||||
|
ERRORS.labels(source="modal").inc()
|
||||||
|
logger.error(
|
||||||
|
"Unhandled exception in modal %r.",
|
||||||
|
self,
|
||||||
|
exc_info=error,
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,122 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""TTL cache for recently ingested messages, enabling uningest on delete."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class CachedMessage:
|
||||||
|
"""Data stored for a recently ingested message.
|
||||||
|
|
||||||
|
:param channel_id: The Discord channel ID the message was ingested into.
|
||||||
|
:param user_id: The Discord user ID of the message author.
|
||||||
|
:param content: The text content of the message, if any.
|
||||||
|
:param embed_texts: Titles and descriptions extracted from embeds.
|
||||||
|
:param image_urls: Image URLs extracted from embeds.
|
||||||
|
"""
|
||||||
|
|
||||||
|
channel_id: int
|
||||||
|
user_id: int
|
||||||
|
content: str | None
|
||||||
|
embed_texts: list[str]
|
||||||
|
image_urls: list[str]
|
||||||
|
|
||||||
|
|
||||||
|
class IngestCache:
|
||||||
|
"""In-memory TTL cache mapping message IDs to their ingested data.
|
||||||
|
|
||||||
|
Entries expire after ``ttl_seconds`` and are evicted by a periodic
|
||||||
|
background task.
|
||||||
|
|
||||||
|
:param ttl_seconds: How long entries remain valid, in seconds.
|
||||||
|
:param cleanup_interval_seconds: How often the background task runs.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__slots__ = ("_cleanup_interval", "_entries", "_task", "_ttl")
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
ttl_seconds: float = 120,
|
||||||
|
cleanup_interval_seconds: float = 30,
|
||||||
|
) -> None:
|
||||||
|
"""Create a new cache with the given TTL and cleanup interval.
|
||||||
|
|
||||||
|
:param ttl_seconds: How long entries remain valid, in seconds.
|
||||||
|
:param cleanup_interval_seconds: How often the background task runs.
|
||||||
|
"""
|
||||||
|
self._ttl = ttl_seconds
|
||||||
|
self._cleanup_interval = cleanup_interval_seconds
|
||||||
|
self._entries: dict[int, tuple[float, CachedMessage]] = {}
|
||||||
|
self._task: asyncio.Task[None] | None = None
|
||||||
|
|
||||||
|
def put(self, message_id: int, entry: CachedMessage) -> None:
|
||||||
|
"""Store a cache entry for a message.
|
||||||
|
|
||||||
|
:param message_id: The Discord message ID.
|
||||||
|
:param entry: The cached message data.
|
||||||
|
"""
|
||||||
|
self._entries[message_id] = (time.monotonic(), entry)
|
||||||
|
|
||||||
|
def pop(self, message_id: int) -> CachedMessage | None:
|
||||||
|
"""Remove and return a cache entry if it exists and has not expired.
|
||||||
|
|
||||||
|
:param message_id: The Discord message ID.
|
||||||
|
:return: The cached message data, or None.
|
||||||
|
"""
|
||||||
|
pair = self._entries.pop(message_id, None)
|
||||||
|
if pair is None:
|
||||||
|
return None
|
||||||
|
stored_at, entry = pair
|
||||||
|
if time.monotonic() - stored_at > self._ttl:
|
||||||
|
return None
|
||||||
|
return entry
|
||||||
|
|
||||||
|
def _cleanup(self) -> None:
|
||||||
|
"""Remove all expired entries from the cache."""
|
||||||
|
now = time.monotonic()
|
||||||
|
expired = [
|
||||||
|
mid
|
||||||
|
for mid, (stored_at, _) in self._entries.items()
|
||||||
|
if now - stored_at > self._ttl
|
||||||
|
]
|
||||||
|
for mid in expired:
|
||||||
|
del self._entries[mid]
|
||||||
|
|
||||||
|
def start(self) -> None:
|
||||||
|
"""Start the periodic background cleanup task."""
|
||||||
|
if self._task is not None:
|
||||||
|
return
|
||||||
|
self._task = asyncio.create_task(
|
||||||
|
self._cleanup_loop(),
|
||||||
|
name="ingest-cache-cleanup",
|
||||||
|
)
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""Stop the periodic background cleanup task and wait for it to finish."""
|
||||||
|
if self._task is not None:
|
||||||
|
self._task.cancel()
|
||||||
|
with contextlib.suppress(asyncio.CancelledError):
|
||||||
|
await self._task
|
||||||
|
self._task = None
|
||||||
|
|
||||||
|
async def _cleanup_loop(self) -> None:
|
||||||
|
"""Run cleanup on a fixed interval until cancelled."""
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(self._cleanup_interval)
|
||||||
|
self._cleanup()
|
||||||
+258
-130
@@ -12,25 +12,42 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Async SQLite database access for Crabstero's persistent storage."""
|
||||||
Provides async SQLite database access for all of Crabstero's persistent storage needs.
|
|
||||||
|
|
||||||
Uses aiosqlite for native async access. All methods are async def.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
from contextlib import asynccontextmanager
|
||||||
import logging
|
from typing import TYPE_CHECKING, NamedTuple, Self
|
||||||
from typing import Self
|
|
||||||
|
|
||||||
import aiosqlite
|
import aiosqlite
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
|
||||||
|
|
||||||
|
class StartWord(NamedTuple):
|
||||||
|
"""A Markov chain starting word row."""
|
||||||
|
|
||||||
|
channel_id: int
|
||||||
|
user_id: int
|
||||||
|
word: str
|
||||||
|
|
||||||
|
|
||||||
|
class Transition(NamedTuple):
|
||||||
|
"""A Markov chain word transition row."""
|
||||||
|
|
||||||
|
channel_id: int
|
||||||
|
user_id: int
|
||||||
|
word: str
|
||||||
|
next_word: str
|
||||||
|
|
||||||
|
|
||||||
|
class ChannelImage(NamedTuple):
|
||||||
|
"""An image URL associated with a channel."""
|
||||||
|
|
||||||
|
channel_id: int
|
||||||
|
user_id: int
|
||||||
|
url: str
|
||||||
|
|
||||||
AUTO_COMMIT_WRITE_THRESHOLD = 100 # Commit after this many DB write operations.
|
|
||||||
AUTO_COMMIT_TIMEOUT_SECONDS = (
|
|
||||||
60.0 # Commit after this many seconds since first uncommitted write.
|
|
||||||
)
|
|
||||||
|
|
||||||
# SQL statements for creating the database schema.
|
# SQL statements for creating the database schema.
|
||||||
_SCHEMA = """
|
_SCHEMA = """
|
||||||
@@ -38,25 +55,32 @@ _SCHEMA = """
|
|||||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||||
CREATE TABLE IF NOT EXISTS markov_start_words (
|
CREATE TABLE IF NOT EXISTS markov_start_words (
|
||||||
channel_id INTEGER NOT NULL,
|
channel_id INTEGER NOT NULL,
|
||||||
|
user_id INTEGER NOT NULL,
|
||||||
word TEXT NOT NULL
|
word TEXT NOT NULL
|
||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS idx_start_channel ON markov_start_words(channel_id, word);
|
CREATE INDEX IF NOT EXISTS idx_start_channel ON markov_start_words(channel_id, word);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_start_user ON markov_start_words(user_id);
|
||||||
|
|
||||||
-- Markov chain word transitions.
|
-- Markov chain word transitions.
|
||||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||||
CREATE TABLE IF NOT EXISTS markov_transitions (
|
CREATE TABLE IF NOT EXISTS markov_transitions (
|
||||||
channel_id INTEGER NOT NULL,
|
channel_id INTEGER NOT NULL,
|
||||||
|
user_id INTEGER NOT NULL,
|
||||||
word TEXT NOT NULL,
|
word TEXT NOT NULL,
|
||||||
next_word TEXT NOT NULL
|
next_word TEXT NOT NULL
|
||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS idx_transitions_channel_word ON markov_transitions(channel_id, word);
|
CREATE INDEX IF NOT EXISTS idx_transitions_channel_word
|
||||||
|
ON markov_transitions(channel_id, word);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_transitions_user ON markov_transitions(user_id);
|
||||||
|
|
||||||
-- Image URLs per channel.
|
-- Image URLs per channel.
|
||||||
CREATE TABLE IF NOT EXISTS channel_images (
|
CREATE TABLE IF NOT EXISTS channel_images (
|
||||||
channel_id INTEGER NOT NULL,
|
channel_id INTEGER NOT NULL,
|
||||||
|
user_id INTEGER NOT NULL,
|
||||||
url TEXT NOT NULL
|
url TEXT NOT NULL
|
||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS idx_images_channel ON channel_images(channel_id);
|
CREATE INDEX IF NOT EXISTS idx_images_channel ON channel_images(channel_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_images_user ON channel_images(user_id);
|
||||||
|
|
||||||
-- Flags for channels, servers, and users.
|
-- Flags for channels, servers, and users.
|
||||||
CREATE TABLE IF NOT EXISTS flags (
|
CREATE TABLE IF NOT EXISTS flags (
|
||||||
@@ -74,36 +98,48 @@ CREATE TABLE IF NOT EXISTS ingested_channels (
|
|||||||
|
|
||||||
|
|
||||||
class Database:
|
class Database:
|
||||||
"""
|
"""Manages all SQLite database operations for Crabstero.
|
||||||
Manages all SQLite database operations for Crabstero.
|
|
||||||
|
|
||||||
Uses aiosqlite for native async access. A single connection is held open for the lifetime of
|
Uses aiosqlite for native async access. A single connection is held
|
||||||
the bot process with WAL mode enabled for concurrent read performance.
|
open for the lifetime of the bot process.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, connection: aiosqlite.Connection) -> None:
|
def __init__(self, connection: aiosqlite.Connection) -> None:
|
||||||
"""
|
"""Initialize the Database wrapper with an already-opened aiosqlite connection.
|
||||||
Initializes the Database wrapper with an already-opened aiosqlite connection.
|
|
||||||
|
|
||||||
:param connection: An open aiosqlite connection.
|
:param connection: An open aiosqlite connection.
|
||||||
"""
|
"""
|
||||||
self._connection = connection
|
self._connection = connection
|
||||||
self._pending_writes = 0
|
self._tx_lock = asyncio.Lock()
|
||||||
self._flush_task: asyncio.Task[None] | None = None
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def _transaction(self) -> AsyncIterator[None]:
|
||||||
|
"""Begin an immediate write transaction.
|
||||||
|
|
||||||
|
Acquires an async lock so that only one transaction runs at a
|
||||||
|
time, then commits on success or rolls back on error.
|
||||||
|
"""
|
||||||
|
async with self._tx_lock:
|
||||||
|
await self._connection.execute("BEGIN IMMEDIATE")
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
except BaseException:
|
||||||
|
await self._connection.rollback()
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
await self._connection.commit()
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
async def connect(cls, path: str) -> Self:
|
async def connect(cls, path: str) -> Self:
|
||||||
"""
|
"""Open a SQLite database, configure it, and create the schema.
|
||||||
Opens a new SQLite database at the given path, configures it for performance, and creates
|
|
||||||
the schema if it does not already exist.
|
Configure the database for performance and create the schema if it
|
||||||
|
does not already exist.
|
||||||
|
|
||||||
:param path: The file path to the SQLite database.
|
:param path: The file path to the SQLite database.
|
||||||
:return: A new Database instance ready for use.
|
:return: A new Database instance ready for use.
|
||||||
"""
|
"""
|
||||||
connection = await aiosqlite.connect(path)
|
connection = await aiosqlite.connect(path, isolation_level=None)
|
||||||
|
|
||||||
# Enable WAL mode for better concurrent read performance.
|
|
||||||
await connection.execute("PRAGMA journal_mode=WAL")
|
|
||||||
|
|
||||||
# Set synchronous to NORMAL for a balance between safety and speed.
|
# Set synchronous to NORMAL for a balance between safety and speed.
|
||||||
await connection.execute("PRAGMA synchronous=NORMAL")
|
await connection.execute("PRAGMA synchronous=NORMAL")
|
||||||
@@ -114,176 +150,254 @@ class Database:
|
|||||||
|
|
||||||
return cls(connection)
|
return cls(connection)
|
||||||
|
|
||||||
async def commit(self) -> None:
|
|
||||||
"""Commits pending writes and resets the flush timer."""
|
|
||||||
await self._connection.commit()
|
|
||||||
self._pending_writes = 0
|
|
||||||
if self._flush_task is not None:
|
|
||||||
self._flush_task.cancel()
|
|
||||||
self._flush_task = None
|
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Cancels the flush timer, commits any pending writes, then closes the database connection."""
|
"""Close the database connection."""
|
||||||
if self._flush_task is not None:
|
|
||||||
self._flush_task.cancel()
|
|
||||||
with contextlib.suppress(asyncio.CancelledError):
|
|
||||||
await self._flush_task
|
|
||||||
self._flush_task = None
|
|
||||||
if self._pending_writes > 0:
|
|
||||||
await self._connection.commit()
|
|
||||||
self._pending_writes = 0
|
|
||||||
await self._connection.close()
|
await self._connection.close()
|
||||||
|
|
||||||
async def add_start_words_batch(self, rows: list[tuple[int, str]]) -> None:
|
async def add_markov_data(
|
||||||
"""
|
self,
|
||||||
Inserts a batch of starting words into the markov_start_words table.
|
start_words: list[StartWord],
|
||||||
|
transitions: list[Transition],
|
||||||
|
) -> None:
|
||||||
|
"""Insert Markov start words and transitions, then commit.
|
||||||
|
|
||||||
:param rows: A list of (channel_id, word) tuples to insert.
|
Both inserts happen in a single transaction.
|
||||||
|
|
||||||
|
:param start_words: Starting word entries to insert.
|
||||||
|
:param transitions: Transition entries to insert.
|
||||||
"""
|
"""
|
||||||
|
if start_words or transitions:
|
||||||
|
async with self._transaction():
|
||||||
|
if start_words:
|
||||||
await self._connection.executemany(
|
await self._connection.executemany(
|
||||||
"INSERT INTO markov_start_words (channel_id, word) VALUES (?, ?)", rows
|
"INSERT INTO markov_start_words"
|
||||||
|
" (channel_id, user_id, word)"
|
||||||
|
" VALUES (?, ?, ?)",
|
||||||
|
start_words,
|
||||||
|
)
|
||||||
|
if transitions:
|
||||||
|
await self._connection.executemany(
|
||||||
|
"INSERT INTO markov_transitions"
|
||||||
|
" (channel_id, user_id, word, next_word)"
|
||||||
|
" VALUES (?, ?, ?, ?)",
|
||||||
|
transitions,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def remove_markov_data(
|
||||||
|
self,
|
||||||
|
start_words: list[StartWord],
|
||||||
|
transitions: list[Transition],
|
||||||
|
) -> None:
|
||||||
|
"""Remove one matching row per entry from the Markov tables.
|
||||||
|
|
||||||
|
Each entry removes at most one duplicate row, preserving remaining
|
||||||
|
frequency weight.
|
||||||
|
|
||||||
|
:param start_words: Starting word entries to remove.
|
||||||
|
:param transitions: Transition entries to remove.
|
||||||
|
"""
|
||||||
|
if start_words or transitions:
|
||||||
|
async with self._transaction():
|
||||||
|
if start_words:
|
||||||
|
await self._connection.executemany(
|
||||||
|
"DELETE FROM markov_start_words"
|
||||||
|
" WHERE rowid = ("
|
||||||
|
" SELECT rowid FROM markov_start_words"
|
||||||
|
" WHERE channel_id = ?"
|
||||||
|
" AND user_id = ?"
|
||||||
|
" AND word = ?"
|
||||||
|
" LIMIT 1"
|
||||||
|
" )",
|
||||||
|
start_words,
|
||||||
|
)
|
||||||
|
if transitions:
|
||||||
|
await self._connection.executemany(
|
||||||
|
"DELETE FROM markov_transitions"
|
||||||
|
" WHERE rowid = ("
|
||||||
|
" SELECT rowid"
|
||||||
|
" FROM markov_transitions"
|
||||||
|
" WHERE channel_id = ?"
|
||||||
|
" AND user_id = ?"
|
||||||
|
" AND word = ?"
|
||||||
|
" AND next_word = ?"
|
||||||
|
" LIMIT 1"
|
||||||
|
" )",
|
||||||
|
transitions,
|
||||||
)
|
)
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def get_random_start_word(self, channel_id: int) -> str | None:
|
async def get_random_start_word(self, channel_id: int) -> str | None:
|
||||||
"""
|
"""Return a random starting word for a channel.
|
||||||
Returns a random starting word for a given channel, weighted by occurrence frequency.
|
|
||||||
|
Weighted by occurrence frequency.
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
:return: A random starting word, or None if no starting words exist for this channel.
|
:return: A random starting word, or None if none exist.
|
||||||
"""
|
"""
|
||||||
async with self._connection.execute(
|
async with self._connection.execute(
|
||||||
"SELECT word FROM markov_start_words WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
"SELECT word FROM markov_start_words"
|
||||||
|
" WHERE channel_id = ?"
|
||||||
|
" ORDER BY RANDOM() LIMIT 1",
|
||||||
(channel_id,),
|
(channel_id,),
|
||||||
) as cursor:
|
) as cursor:
|
||||||
row = await cursor.fetchone()
|
row = await cursor.fetchone()
|
||||||
return row[0] if row else None
|
return row[0] if row else None
|
||||||
|
|
||||||
async def add_transitions_batch(self, rows: list[tuple[int, str, str]]) -> None:
|
|
||||||
"""
|
|
||||||
Inserts a batch of word transitions into the markov_transitions table.
|
|
||||||
|
|
||||||
:param rows: A list of (channel_id, word, next_word) tuples to insert.
|
|
||||||
"""
|
|
||||||
await self._connection.executemany(
|
|
||||||
"INSERT INTO markov_transitions (channel_id, word, next_word) VALUES (?, ?, ?)",
|
|
||||||
rows,
|
|
||||||
)
|
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def get_random_next_word(self, channel_id: int, word: str) -> str | None:
|
async def get_random_next_word(self, channel_id: int, word: str) -> str | None:
|
||||||
"""
|
"""Return a random next word for a given word in a channel.
|
||||||
Returns a random next word for a given word in a given channel, weighted by occurrence
|
|
||||||
frequency.
|
Weighted by occurrence frequency.
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
:param word: The current word to find a transition for.
|
:param word: The current word to find a transition for.
|
||||||
:return: A random next word, or None if no transitions exist.
|
:return: A random next word, or None if no transitions exist.
|
||||||
"""
|
"""
|
||||||
async with self._connection.execute(
|
async with self._connection.execute(
|
||||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? ORDER BY RANDOM() LIMIT 1",
|
"SELECT next_word FROM markov_transitions"
|
||||||
|
" WHERE channel_id = ? AND word = ?"
|
||||||
|
" ORDER BY RANDOM() LIMIT 1",
|
||||||
(channel_id, word),
|
(channel_id, word),
|
||||||
) as cursor:
|
) as cursor:
|
||||||
row = await cursor.fetchone()
|
row = await cursor.fetchone()
|
||||||
return row[0] if row else None
|
return row[0] if row else None
|
||||||
|
|
||||||
async def get_random_completing_next_word(
|
async def get_random_completing_next_word(
|
||||||
self, channel_id: int, word: str
|
self,
|
||||||
|
channel_id: int,
|
||||||
|
word: str,
|
||||||
) -> str | None:
|
) -> str | None:
|
||||||
"""
|
"""Return a random sentence-ending next word for a given word in a channel.
|
||||||
Returns a random next word that ends a sentence for a given word in a given channel,
|
|
||||||
weighted by occurrence frequency.
|
|
||||||
|
|
||||||
A completing word is one whose last character is '.', '!', '?', or '§'.
|
Weighted by occurrence frequency. A completing word is one whose last
|
||||||
|
character is '.', '!', '?', or '§'.
|
||||||
|
|
||||||
|
The set of terminators must match ``markov._TERMINATORS``.
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
:param word: The current word to find a completing transition for.
|
:param word: The current word to find a completing transition for.
|
||||||
:return: A random completing next word, or None if no completing transitions exist.
|
:return: A random completing next word, or None if none exist.
|
||||||
"""
|
"""
|
||||||
async with self._connection.execute(
|
async with self._connection.execute(
|
||||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? AND SUBSTR(next_word, -1, 1) IN ('.', '!', '?', '§') ORDER BY RANDOM() LIMIT 1",
|
"SELECT next_word FROM markov_transitions"
|
||||||
|
" WHERE channel_id = ? AND word = ?"
|
||||||
|
" AND SUBSTR(next_word, -1, 1)"
|
||||||
|
" IN ('.', '!', '?', '§')"
|
||||||
|
" ORDER BY RANDOM() LIMIT 1",
|
||||||
(channel_id, word),
|
(channel_id, word),
|
||||||
) as cursor:
|
) as cursor:
|
||||||
row = await cursor.fetchone()
|
row = await cursor.fetchone()
|
||||||
return row[0] if row else None
|
return row[0] if row else None
|
||||||
|
|
||||||
async def add_image(self, channel_id: int, url: str) -> None:
|
async def add_images(self, images: list[ChannelImage]) -> None:
|
||||||
"""
|
"""Store image URLs for a given channel.
|
||||||
Stores an image URL for a given channel.
|
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param images: Image entries to insert.
|
||||||
:param url: The image URL to store.
|
|
||||||
"""
|
"""
|
||||||
await self._connection.execute(
|
if images:
|
||||||
"INSERT INTO channel_images (channel_id, url) VALUES (?, ?)",
|
async with self._transaction():
|
||||||
(channel_id, url),
|
await self._connection.executemany(
|
||||||
|
"INSERT INTO channel_images"
|
||||||
|
" (channel_id, user_id, url)"
|
||||||
|
" VALUES (?, ?, ?)",
|
||||||
|
images,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def remove_images(self, images: list[ChannelImage]) -> None:
|
||||||
|
"""Remove one matching row per entry from the images table.
|
||||||
|
|
||||||
|
Each entry removes at most one duplicate row, preserving remaining
|
||||||
|
frequency weight.
|
||||||
|
|
||||||
|
:param images: Image entries to remove.
|
||||||
|
"""
|
||||||
|
if images:
|
||||||
|
async with self._transaction():
|
||||||
|
await self._connection.executemany(
|
||||||
|
"DELETE FROM channel_images"
|
||||||
|
" WHERE rowid = ("
|
||||||
|
" SELECT rowid FROM channel_images"
|
||||||
|
" WHERE channel_id = ?"
|
||||||
|
" AND user_id = ?"
|
||||||
|
" AND url = ?"
|
||||||
|
" LIMIT 1"
|
||||||
|
" )",
|
||||||
|
images,
|
||||||
)
|
)
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def get_random_image(self, channel_id: int) -> str | None:
|
async def get_random_image(self, channel_id: int) -> str | None:
|
||||||
"""
|
"""Return a random image URL for a given channel.
|
||||||
Returns a random image URL for a given channel.
|
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
:return: A random image URL, or None if no images exist for this channel.
|
:return: A random image URL, or None if none exist.
|
||||||
"""
|
"""
|
||||||
async with self._connection.execute(
|
async with self._connection.execute(
|
||||||
"SELECT url FROM channel_images WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
"SELECT url FROM channel_images"
|
||||||
|
" WHERE channel_id = ?"
|
||||||
|
" ORDER BY RANDOM() LIMIT 1",
|
||||||
(channel_id,),
|
(channel_id,),
|
||||||
) as cursor:
|
) as cursor:
|
||||||
row = await cursor.fetchone()
|
row = await cursor.fetchone()
|
||||||
return row[0] if row else None
|
return row[0] if row else None
|
||||||
|
|
||||||
async def set_flag(self, entity_type: str, entity_id: str, flag_name: str) -> None:
|
async def set_flag(self, entity_type: str, entity_id: str, flag_name: str) -> None:
|
||||||
"""
|
"""Set a flag on a given entity. If the flag is already set, this is a no-op.
|
||||||
Sets a flag on a given entity. If the flag is already set, this is a no-op.
|
|
||||||
|
|
||||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
:param entity_type: The entity type ("channel", "server", "user").
|
||||||
:param entity_id: The Discord ID of the entity.
|
:param entity_id: The Discord ID of the entity.
|
||||||
:param flag_name: The name of the flag to set.
|
:param flag_name: The name of the flag to set.
|
||||||
"""
|
"""
|
||||||
|
async with self._transaction():
|
||||||
await self._connection.execute(
|
await self._connection.execute(
|
||||||
"INSERT OR IGNORE INTO flags (entity_type, entity_id, flag_name) VALUES (?, ?, ?)",
|
"INSERT OR IGNORE INTO flags"
|
||||||
|
" (entity_type, entity_id, flag_name)"
|
||||||
|
" VALUES (?, ?, ?)",
|
||||||
(entity_type, entity_id, flag_name),
|
(entity_type, entity_id, flag_name),
|
||||||
)
|
)
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def clear_flag(
|
async def clear_flag(
|
||||||
self, entity_type: str, entity_id: str, flag_name: str
|
self,
|
||||||
|
entity_type: str,
|
||||||
|
entity_id: str,
|
||||||
|
flag_name: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Clear a flag on a given entity. If the flag is not set, this is a no-op.
|
||||||
Clears a flag on a given entity. If the flag is not set, this is a no-op.
|
|
||||||
|
|
||||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
:param entity_type: The entity type ("channel", "server", "user").
|
||||||
:param entity_id: The Discord ID of the entity.
|
:param entity_id: The Discord ID of the entity.
|
||||||
:param flag_name: The name of the flag to clear.
|
:param flag_name: The name of the flag to clear.
|
||||||
"""
|
"""
|
||||||
|
async with self._transaction():
|
||||||
await self._connection.execute(
|
await self._connection.execute(
|
||||||
"DELETE FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
"DELETE FROM flags"
|
||||||
|
" WHERE entity_type = ?"
|
||||||
|
" AND entity_id = ?"
|
||||||
|
" AND flag_name = ?",
|
||||||
(entity_type, entity_id, flag_name),
|
(entity_type, entity_id, flag_name),
|
||||||
)
|
)
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def is_flag_set(
|
async def is_flag_set(
|
||||||
self, entity_type: str, entity_id: str, flag_name: str
|
self,
|
||||||
|
entity_type: str,
|
||||||
|
entity_id: str,
|
||||||
|
flag_name: str,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""
|
"""Check whether a flag is set on a given entity.
|
||||||
Checks whether a flag is set on a given entity.
|
|
||||||
|
|
||||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
:param entity_type: The entity type ("channel", "server", "user").
|
||||||
:param entity_id: The Discord ID of the entity.
|
:param entity_id: The Discord ID of the entity.
|
||||||
:param flag_name: The name of the flag to check.
|
:param flag_name: The name of the flag to check.
|
||||||
:return: True if the flag is set, False otherwise.
|
:return: True if the flag is set, False otherwise.
|
||||||
"""
|
"""
|
||||||
async with self._connection.execute(
|
async with self._connection.execute(
|
||||||
"SELECT 1 FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
"SELECT 1 FROM flags"
|
||||||
|
" WHERE entity_type = ?"
|
||||||
|
" AND entity_id = ?"
|
||||||
|
" AND flag_name = ?",
|
||||||
(entity_type, entity_id, flag_name),
|
(entity_type, entity_id, flag_name),
|
||||||
) as cursor:
|
) as cursor:
|
||||||
return await cursor.fetchone() is not None
|
return await cursor.fetchone() is not None
|
||||||
|
|
||||||
async def is_channel_ingested(self, channel_id: int) -> bool:
|
async def is_channel_ingested(self, channel_id: int) -> bool:
|
||||||
"""
|
"""Check whether a channel has already been bulk-ingested.
|
||||||
Checks whether a channel has already been bulk-ingested.
|
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
:return: True if the channel has been ingested, False otherwise.
|
:return: True if the channel has been ingested, False otherwise.
|
||||||
@@ -295,32 +409,46 @@ class Database:
|
|||||||
return await cursor.fetchone() is not None
|
return await cursor.fetchone() is not None
|
||||||
|
|
||||||
async def mark_channel_ingested(self, channel_id: int) -> None:
|
async def mark_channel_ingested(self, channel_id: int) -> None:
|
||||||
"""
|
"""Mark a channel as having been bulk-ingested.
|
||||||
Marks a channel as having been bulk-ingested.
|
|
||||||
|
|
||||||
:param channel_id: The Discord channel ID.
|
:param channel_id: The Discord channel ID.
|
||||||
"""
|
"""
|
||||||
|
async with self._transaction():
|
||||||
await self._connection.execute(
|
await self._connection.execute(
|
||||||
"INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)",
|
"INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)",
|
||||||
(channel_id,),
|
(channel_id,),
|
||||||
)
|
)
|
||||||
await self._maybe_commit()
|
|
||||||
|
|
||||||
async def _maybe_commit(self) -> None:
|
async def forget_user(self, user_id: int, no_ingest_flag: str) -> None:
|
||||||
"""Tracks a pending write. Commits if the threshold is reached, otherwise starts a flush timer."""
|
"""Delete all user data, clear flags, and set noIngest atomically.
|
||||||
self._pending_writes += 1
|
|
||||||
if self._pending_writes >= AUTO_COMMIT_WRITE_THRESHOLD:
|
|
||||||
await self.commit()
|
|
||||||
elif self._flush_task is None:
|
|
||||||
self._flush_task = asyncio.create_task(self._flush_after_timeout())
|
|
||||||
|
|
||||||
async def _flush_after_timeout(self) -> None:
|
Sets the noIngest flag, clears all other user flags, and deletes
|
||||||
"""Background task that commits after the timeout elapses."""
|
all Markov and image data in a single transaction.
|
||||||
try:
|
|
||||||
await asyncio.sleep(AUTO_COMMIT_TIMEOUT_SECONDS)
|
:param user_id: The Discord user ID to forget.
|
||||||
if self._pending_writes > 0:
|
:param no_ingest_flag: The flag name for noIngest.
|
||||||
await self.commit()
|
"""
|
||||||
except asyncio.CancelledError:
|
entity_id = str(user_id)
|
||||||
raise
|
async with self._transaction():
|
||||||
except Exception:
|
await self._connection.execute(
|
||||||
logger.exception("Error in database flush timer.")
|
"INSERT OR IGNORE INTO flags (entity_type, entity_id, flag_name)"
|
||||||
|
" VALUES ('user', ?, ?)",
|
||||||
|
(entity_id, no_ingest_flag),
|
||||||
|
)
|
||||||
|
await self._connection.execute(
|
||||||
|
"DELETE FROM flags WHERE entity_type = 'user'"
|
||||||
|
" AND entity_id = ? AND flag_name != ?",
|
||||||
|
(entity_id, no_ingest_flag),
|
||||||
|
)
|
||||||
|
await self._connection.execute(
|
||||||
|
"DELETE FROM markov_start_words WHERE user_id = ?",
|
||||||
|
(user_id,),
|
||||||
|
)
|
||||||
|
await self._connection.execute(
|
||||||
|
"DELETE FROM markov_transitions WHERE user_id = ?",
|
||||||
|
(user_id,),
|
||||||
|
)
|
||||||
|
await self._connection.execute(
|
||||||
|
"DELETE FROM channel_images WHERE user_id = ?",
|
||||||
|
(user_id,),
|
||||||
|
)
|
||||||
|
|||||||
+10
-15
@@ -12,8 +12,7 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Provides async convenience wrappers around the database flag methods.
|
||||||
Provides async convenience wrappers around the database flag methods.
|
|
||||||
|
|
||||||
Accepts discord.py objects or raw integer IDs and translates them into the
|
Accepts discord.py objects or raw integer IDs and translates them into the
|
||||||
entity_type/entity_id pairs used by the database layer.
|
entity_type/entity_id pairs used by the database layer.
|
||||||
@@ -28,7 +27,7 @@ if TYPE_CHECKING:
|
|||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
|
||||||
|
|
||||||
class Flag(enum.Enum):
|
class Flag(enum.StrEnum):
|
||||||
"""Enum of all supported flag names."""
|
"""Enum of all supported flag names."""
|
||||||
|
|
||||||
NO_REPLY = "noReply"
|
NO_REPLY = "noReply"
|
||||||
@@ -36,7 +35,7 @@ class Flag(enum.Enum):
|
|||||||
ALLOW_PINGS = "allowPings"
|
ALLOW_PINGS = "allowPings"
|
||||||
|
|
||||||
|
|
||||||
class EntityType(enum.Enum):
|
class EntityType(enum.StrEnum):
|
||||||
"""Enum of entity types that can have flags."""
|
"""Enum of entity types that can have flags."""
|
||||||
|
|
||||||
CHANNEL = "channel"
|
CHANNEL = "channel"
|
||||||
@@ -45,8 +44,7 @@ class EntityType(enum.Enum):
|
|||||||
|
|
||||||
|
|
||||||
def _entity_id(entity: discord.abc.Snowflake | int) -> str:
|
def _entity_id(entity: discord.abc.Snowflake | int) -> str:
|
||||||
"""
|
"""Extract a string entity ID from a Discord object or raw integer ID.
|
||||||
Extracts a string entity ID from a Discord object or raw integer ID.
|
|
||||||
|
|
||||||
:param entity: A Discord entity or raw integer ID.
|
:param entity: A Discord entity or raw integer ID.
|
||||||
:return: The entity ID as a string.
|
:return: The entity ID as a string.
|
||||||
@@ -60,15 +58,14 @@ async def set_flag(
|
|||||||
entity_type: EntityType,
|
entity_type: EntityType,
|
||||||
flag: Flag,
|
flag: Flag,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Set a flag on a given entity.
|
||||||
Sets a flag on a given entity.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param entity: The Discord entity or raw integer ID to set the flag on.
|
:param entity: The Discord entity or raw integer ID to set the flag on.
|
||||||
:param entity_type: The type of the entity.
|
:param entity_type: The type of the entity.
|
||||||
:param flag: The flag to set.
|
:param flag: The flag to set.
|
||||||
"""
|
"""
|
||||||
await db.set_flag(entity_type.value, _entity_id(entity), flag.value)
|
await db.set_flag(entity_type, _entity_id(entity), flag)
|
||||||
|
|
||||||
|
|
||||||
async def clear_flag(
|
async def clear_flag(
|
||||||
@@ -77,15 +74,14 @@ async def clear_flag(
|
|||||||
entity_type: EntityType,
|
entity_type: EntityType,
|
||||||
flag: Flag,
|
flag: Flag,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Clear a flag on a given entity.
|
||||||
Clears a flag on a given entity.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param entity: The Discord entity or raw integer ID to clear the flag on.
|
:param entity: The Discord entity or raw integer ID to clear the flag on.
|
||||||
:param entity_type: The type of the entity.
|
:param entity_type: The type of the entity.
|
||||||
:param flag: The flag to clear.
|
:param flag: The flag to clear.
|
||||||
"""
|
"""
|
||||||
await db.clear_flag(entity_type.value, _entity_id(entity), flag.value)
|
await db.clear_flag(entity_type, _entity_id(entity), flag)
|
||||||
|
|
||||||
|
|
||||||
async def is_flag_set(
|
async def is_flag_set(
|
||||||
@@ -94,8 +90,7 @@ async def is_flag_set(
|
|||||||
entity_type: EntityType,
|
entity_type: EntityType,
|
||||||
flag: Flag,
|
flag: Flag,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""
|
"""Check whether a flag is set on a given entity.
|
||||||
Checks whether a flag is set on a given entity.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param entity: The Discord entity or raw integer ID to check.
|
:param entity: The Discord entity or raw integer ID to check.
|
||||||
@@ -103,4 +98,4 @@ async def is_flag_set(
|
|||||||
:param flag: The flag to check for.
|
:param flag: The flag to check for.
|
||||||
:return: True if the flag is set, False otherwise.
|
:return: True if the flag is set, False otherwise.
|
||||||
"""
|
"""
|
||||||
return await db.is_flag_set(entity_type.value, _entity_id(entity), flag.value)
|
return await db.is_flag_set(entity_type, _entity_id(entity), flag)
|
||||||
|
|||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Listener cogs for Discord events."""
|
||||||
|
|||||||
@@ -12,36 +12,129 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Handles responding to interactions.
|
||||||
Handles responding to interactions.
|
|
||||||
|
|
||||||
Provides the /pingme slash command as a Cog with an app command.
|
Provides the /pingme and /forgetme slash commands as a Cog with app commands.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING, override
|
||||||
|
|
||||||
import discord
|
import discord
|
||||||
from discord import app_commands
|
from discord import app_commands
|
||||||
from discord.ext import commands
|
from discord.ext import commands
|
||||||
|
|
||||||
from crabstero import flags
|
from crabstero import flags, metrics
|
||||||
|
from crabstero.bot import TrackedView
|
||||||
from crabstero.flags import EntityType, Flag
|
from crabstero.flags import EntityType, Flag
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crabstero.bot import Crabstero
|
from crabstero.bot import Crabstero
|
||||||
|
from crabstero.database import Database
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ForgetMeView(TrackedView):
|
||||||
|
"""Confirmation view with Confirm/Cancel buttons for /forgetme."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
user_id: int,
|
||||||
|
db: Database,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
) -> None:
|
||||||
|
"""Create a new ForgetMeView.
|
||||||
|
|
||||||
|
:param user_id: The Discord user ID who invoked the command.
|
||||||
|
:param db: The database instance.
|
||||||
|
:param interaction: The original interaction for editing on timeout.
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
self._user_id = user_id
|
||||||
|
self._db = db
|
||||||
|
self._responded = False
|
||||||
|
self._interaction = interaction
|
||||||
|
|
||||||
|
@override
|
||||||
|
async def interaction_check(self, interaction: discord.Interaction) -> bool:
|
||||||
|
"""Only allow the invoking user to interact with the buttons.
|
||||||
|
|
||||||
|
:param interaction: The interaction event.
|
||||||
|
:return: True if the user matches, False otherwise.
|
||||||
|
"""
|
||||||
|
return interaction.user.id == self._user_id
|
||||||
|
|
||||||
|
@discord.ui.button(label="Confirm", style=discord.ButtonStyle.danger)
|
||||||
|
async def confirm(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
_button: discord.ui.Button[ForgetMeView],
|
||||||
|
) -> None:
|
||||||
|
"""Delete all user data, clear flags, and set noIngest.
|
||||||
|
|
||||||
|
:param interaction: The interaction event.
|
||||||
|
:param _button: The button that was pressed.
|
||||||
|
"""
|
||||||
|
self._responded = True
|
||||||
|
await self._db.forget_user(self._user_id, Flag.NO_INGEST)
|
||||||
|
metrics.FORGETME.labels(outcome="completed").inc()
|
||||||
|
await interaction.response.edit_message(
|
||||||
|
content=(
|
||||||
|
"I have deleted your data and I will not use"
|
||||||
|
" your messages going forward."
|
||||||
|
),
|
||||||
|
view=None,
|
||||||
|
)
|
||||||
|
self.stop()
|
||||||
|
|
||||||
|
@discord.ui.button(label="Cancel", style=discord.ButtonStyle.secondary)
|
||||||
|
async def cancel(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
_button: discord.ui.Button[ForgetMeView],
|
||||||
|
) -> None:
|
||||||
|
"""Cancel the /forgetme action.
|
||||||
|
|
||||||
|
:param interaction: The interaction event.
|
||||||
|
:param _button: The button that was pressed.
|
||||||
|
"""
|
||||||
|
self._responded = True
|
||||||
|
metrics.FORGETME.labels(outcome="cancelled").inc()
|
||||||
|
await interaction.response.edit_message(
|
||||||
|
content="Action cancelled. I have not modified your data.",
|
||||||
|
view=None,
|
||||||
|
)
|
||||||
|
self.stop()
|
||||||
|
|
||||||
|
@override
|
||||||
|
async def on_timeout(self) -> None:
|
||||||
|
"""Remove buttons and increment timeout metric."""
|
||||||
|
if self._responded:
|
||||||
|
return
|
||||||
|
metrics.FORGETME.labels(outcome="timeout").inc()
|
||||||
|
try:
|
||||||
|
await self._interaction.edit_original_response(
|
||||||
|
content="This timed out. Run `/forgetme` again if you still want to.",
|
||||||
|
view=None,
|
||||||
|
)
|
||||||
|
except discord.NotFound:
|
||||||
|
logger.debug(
|
||||||
|
"Original /forgetme response for user %s was already deleted.",
|
||||||
|
self._user_id,
|
||||||
|
)
|
||||||
|
except discord.HTTPException:
|
||||||
|
logger.debug(
|
||||||
|
"Failed to edit timed-out /forgetme response for user %s.",
|
||||||
|
self._user_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class InteractionCog(commands.Cog):
|
class InteractionCog(commands.Cog):
|
||||||
"""
|
"""Cog for handling slash command interactions."""
|
||||||
Cog for handling slash command interactions.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, bot: Crabstero) -> None:
|
def __init__(self, bot: Crabstero) -> None:
|
||||||
"""
|
"""Create a new interaction handler cog.
|
||||||
Creates a new interaction handler cog.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
@@ -49,48 +142,83 @@ class InteractionCog(commands.Cog):
|
|||||||
|
|
||||||
@app_commands.command(
|
@app_commands.command(
|
||||||
name="pingme",
|
name="pingme",
|
||||||
description="Opts into (or back out of) receiving pings for generated message which mention you.",
|
description=(
|
||||||
|
"Toggle whether you receive pings for generated messages that mention you."
|
||||||
|
),
|
||||||
)
|
)
|
||||||
async def pingme(self, interaction: discord.Interaction) -> None:
|
async def pingme(self, interaction: discord.Interaction) -> None:
|
||||||
"""
|
"""Toggle the allowPings flag for the user who ran the command.
|
||||||
Toggles the allowPings flag for the user who ran the command.
|
|
||||||
|
|
||||||
:param interaction: The interaction event.
|
:param interaction: The interaction event.
|
||||||
"""
|
"""
|
||||||
# Wrap in try/except to always send a response even if the database fails.
|
|
||||||
try:
|
|
||||||
if await flags.is_flag_set(
|
if await flags.is_flag_set(
|
||||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
self.bot.db,
|
||||||
|
interaction.user,
|
||||||
|
EntityType.USER,
|
||||||
|
Flag.ALLOW_PINGS,
|
||||||
):
|
):
|
||||||
await flags.clear_flag(
|
await flags.clear_flag(
|
||||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
self.bot.db,
|
||||||
|
interaction.user,
|
||||||
|
EntityType.USER,
|
||||||
|
Flag.ALLOW_PINGS,
|
||||||
)
|
)
|
||||||
|
metrics.PINGME.labels(outcome="opted_out").inc()
|
||||||
await interaction.response.send_message(
|
await interaction.response.send_message(
|
||||||
"I will no longer ping you for messages which mention you. If you decide to opt back in, run `/pingme` any time.",
|
"I will no longer ping you for messages which"
|
||||||
|
" mention you. If you decide to opt back in,"
|
||||||
|
" run `/pingme` any time.",
|
||||||
ephemeral=True,
|
ephemeral=True,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
await flags.set_flag(
|
await flags.set_flag(
|
||||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
self.bot.db,
|
||||||
|
interaction.user,
|
||||||
|
EntityType.USER,
|
||||||
|
Flag.ALLOW_PINGS,
|
||||||
)
|
)
|
||||||
|
metrics.PINGME.labels(outcome="opted_in").inc()
|
||||||
await interaction.response.send_message(
|
await interaction.response.send_message(
|
||||||
"I will now ping you for messages which mention you. If you change your mind, run `/pingme` any time.",
|
"I will now ping you for messages which mention"
|
||||||
|
" you. If you change your mind, run"
|
||||||
|
" `/pingme` any time.",
|
||||||
ephemeral=True,
|
ephemeral=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
except Exception:
|
@app_commands.command(
|
||||||
logger.exception(
|
name="forgetme",
|
||||||
"An exception occurred while attempting to execute slash command responder for command 'pingme'."
|
description="Delete all your data and stop your messages from being used.",
|
||||||
)
|
)
|
||||||
|
async def forgetme(self, interaction: discord.Interaction) -> None:
|
||||||
|
"""Delete all user data and set the noIngest flag.
|
||||||
|
|
||||||
|
:param interaction: The interaction event.
|
||||||
|
"""
|
||||||
|
if await flags.is_flag_set(
|
||||||
|
self.bot.db,
|
||||||
|
interaction.user,
|
||||||
|
EntityType.USER,
|
||||||
|
Flag.NO_INGEST,
|
||||||
|
):
|
||||||
|
metrics.FORGETME.labels(outcome="already_forgotten").inc()
|
||||||
await interaction.response.send_message(
|
await interaction.response.send_message(
|
||||||
"An error occurred while executing your command. Please try again later.",
|
"I have already removed your data and I am not using your messages.",
|
||||||
|
ephemeral=True,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
view = ForgetMeView(interaction.user.id, self.bot.db, interaction)
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"I will delete everything I have learned from your"
|
||||||
|
" messages and I will not use them going forward."
|
||||||
|
" Would you like to proceed?",
|
||||||
|
view=view,
|
||||||
ephemeral=True,
|
ephemeral=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def setup(bot: Crabstero) -> None:
|
async def setup(bot: Crabstero) -> None:
|
||||||
"""
|
"""Add the InteractionCog to the bot.
|
||||||
Adds the InteractionCog to the bot.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -12,10 +12,10 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Handles message creation and deletion events.
|
||||||
Handles created messages.
|
|
||||||
|
|
||||||
Responds to mentions and ingests normal text messages as a Cog with an on_message listener.
|
Responds to mentions, ingests normal text messages, and reverses
|
||||||
|
ingestion when messages are deleted.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
@@ -23,20 +23,18 @@ from typing import TYPE_CHECKING
|
|||||||
import discord
|
import discord
|
||||||
from discord.ext import commands
|
from discord.ext import commands
|
||||||
|
|
||||||
from crabstero.messages import ingest_message, reply_to_message
|
from crabstero import metrics
|
||||||
|
from crabstero.messages import ingest_message, reply_to_message, uningest_message
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crabstero.bot import Crabstero
|
from crabstero.bot import Crabstero
|
||||||
|
|
||||||
|
|
||||||
class MessageCog(commands.Cog):
|
class MessageCog(commands.Cog):
|
||||||
"""
|
"""Cog for handling message creation events."""
|
||||||
Cog for handling message creation events.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, bot: Crabstero) -> None:
|
def __init__(self, bot: Crabstero) -> None:
|
||||||
"""
|
"""Create a new message creation handler cog.
|
||||||
Creates a new message creation handler cog.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
@@ -44,35 +42,69 @@ class MessageCog(commands.Cog):
|
|||||||
|
|
||||||
@commands.Cog.listener()
|
@commands.Cog.listener()
|
||||||
async def on_message(self, message: discord.Message) -> None:
|
async def on_message(self, message: discord.Message) -> None:
|
||||||
"""
|
"""Respond to mentions and ingest normal text messages.
|
||||||
Responds to mentions and ingests normal text messages.
|
|
||||||
|
|
||||||
:param message: The message event.
|
:param message: The message event.
|
||||||
"""
|
"""
|
||||||
channel = message.channel
|
channel = message.channel
|
||||||
|
|
||||||
if not isinstance(
|
if not isinstance(
|
||||||
channel, (discord.TextChannel, discord.Thread, discord.VoiceChannel)
|
channel,
|
||||||
|
(discord.TextChannel, discord.Thread, discord.VoiceChannel),
|
||||||
):
|
):
|
||||||
return
|
return
|
||||||
|
|
||||||
if message.author.bot or message.author == self.bot.user:
|
if message.author.bot or message.author == self.bot.user:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
metrics.MESSAGES_PROCESSED.inc()
|
||||||
|
|
||||||
|
if self.bot.ingest_only:
|
||||||
|
# In ingest-only mode, never reply — only ingest eligible messages.
|
||||||
|
if message.type == discord.MessageType.default and not isinstance(
|
||||||
|
channel,
|
||||||
|
discord.Thread,
|
||||||
|
):
|
||||||
|
await ingest_message(self.bot.db, message, self.bot.ingest_cache)
|
||||||
|
return
|
||||||
|
|
||||||
if self.bot.user in message.mentions:
|
if self.bot.user in message.mentions:
|
||||||
await reply_to_message(self.bot.db, message)
|
await reply_to_message(self.bot.db, message)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Only ingest non-thread messages; threads share their parent channel's chain.
|
# Only ingest non-thread messages; threads share their parent channel's chain.
|
||||||
if message.type == discord.MessageType.default and not isinstance(
|
if message.type == discord.MessageType.default and not isinstance(
|
||||||
channel, discord.Thread
|
channel,
|
||||||
|
discord.Thread,
|
||||||
):
|
):
|
||||||
await ingest_message(self.bot.db, message)
|
await ingest_message(self.bot.db, message, self.bot.ingest_cache)
|
||||||
|
|
||||||
|
@commands.Cog.listener()
|
||||||
|
async def on_raw_message_delete(
|
||||||
|
self,
|
||||||
|
payload: discord.RawMessageDeleteEvent,
|
||||||
|
) -> None:
|
||||||
|
"""Uningest a recently ingested message when it is deleted.
|
||||||
|
|
||||||
|
:param payload: The raw message delete event.
|
||||||
|
"""
|
||||||
|
await uningest_message(self.bot.db, self.bot.ingest_cache, payload.message_id)
|
||||||
|
|
||||||
|
@commands.Cog.listener()
|
||||||
|
async def on_raw_bulk_message_delete(
|
||||||
|
self,
|
||||||
|
payload: discord.RawBulkMessageDeleteEvent,
|
||||||
|
) -> None:
|
||||||
|
"""Uningest recently ingested messages when they are bulk-deleted.
|
||||||
|
|
||||||
|
:param payload: The raw bulk message delete event.
|
||||||
|
"""
|
||||||
|
for message_id in payload.message_ids:
|
||||||
|
await uningest_message(self.bot.db, self.bot.ingest_cache, message_id)
|
||||||
|
|
||||||
|
|
||||||
async def setup(bot: Crabstero) -> None:
|
async def setup(bot: Crabstero) -> None:
|
||||||
"""
|
"""Add the MessageCog to the bot.
|
||||||
Adds the MessageCog to the bot.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -12,11 +12,11 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Handles server-level events which trigger channel history ingestion.
|
||||||
Handles server-level events which trigger channel history ingestion.
|
|
||||||
|
|
||||||
All of these events share the same outcome: queuing all textable channels for message history
|
All of these events share the same outcome: queuing all textable channels
|
||||||
ingestion when something changes that may grant new read permissions.
|
for message history ingestion when something changes that may grant new
|
||||||
|
read permissions.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import contextlib
|
import contextlib
|
||||||
@@ -35,13 +35,10 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
class ServerEventsCog(commands.Cog):
|
class ServerEventsCog(commands.Cog):
|
||||||
"""
|
"""Cog for handling server-level events that trigger channel history ingestion."""
|
||||||
Cog for handling server-level events that trigger channel history ingestion.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, bot: Crabstero) -> None:
|
def __init__(self, bot: Crabstero) -> None:
|
||||||
"""
|
"""Create a new server events handler cog.
|
||||||
Creates a new server events handler cog.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
@@ -49,29 +46,30 @@ class ServerEventsCog(commands.Cog):
|
|||||||
|
|
||||||
@commands.Cog.listener()
|
@commands.Cog.listener()
|
||||||
async def on_guild_join(self, guild: discord.Guild) -> None:
|
async def on_guild_join(self, guild: discord.Guild) -> None:
|
||||||
"""
|
"""Queue all text channels for ingestion when joining a new server.
|
||||||
Queues all text channels for message history ingestion when joining a new server.
|
|
||||||
Also logs the join and sends an embed to the bot owner.
|
Also log the join and send an embed to the bot owner.
|
||||||
|
|
||||||
:param guild: The guild that was joined.
|
:param guild: The guild that was joined.
|
||||||
"""
|
"""
|
||||||
logger.info('Joined new server! (Name: "%s" | ID: %s)', guild.name, guild.id)
|
logger.info('Joined new server! (Name: "%s" | ID: %s)', guild.name, guild.id)
|
||||||
|
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
title="Joined New Server", color=discord.Color.from_rgb(255, 165, 0)
|
title="Joined New Server",
|
||||||
|
color=discord.Color.from_rgb(255, 165, 0),
|
||||||
)
|
)
|
||||||
embed.add_field(
|
embed.add_field(
|
||||||
name=f"{guild.name} ({guild.id})",
|
name=f"{guild.name} ({guild.id})",
|
||||||
value=f"{guild.member_count} members",
|
value=f"{guild.member_count} members",
|
||||||
)
|
)
|
||||||
|
|
||||||
if guild.icon:
|
if guild.icon is not None:
|
||||||
embed.set_image(url=guild.icon.url)
|
embed.set_image(url=guild.icon.url)
|
||||||
|
|
||||||
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
|
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
|
||||||
|
|
||||||
app_info = await self.bot.application_info()
|
app_info = await self.bot.application_info()
|
||||||
if app_info.owner:
|
if app_info.owner is not None:
|
||||||
with contextlib.suppress(discord.HTTPException):
|
with contextlib.suppress(discord.HTTPException):
|
||||||
await app_info.owner.send(embed=embed)
|
await app_info.owner.send(embed=embed)
|
||||||
|
|
||||||
@@ -79,8 +77,7 @@ class ServerEventsCog(commands.Cog):
|
|||||||
|
|
||||||
@commands.Cog.listener()
|
@commands.Cog.listener()
|
||||||
async def on_guild_available(self, guild: discord.Guild) -> None:
|
async def on_guild_available(self, guild: discord.Guild) -> None:
|
||||||
"""
|
"""Queue all textable channels for ingestion when a server becomes available.
|
||||||
Queues all textable channels for message history ingestion when a server becomes available.
|
|
||||||
|
|
||||||
:param guild: The guild that became available.
|
:param guild: The guild that became available.
|
||||||
"""
|
"""
|
||||||
@@ -88,11 +85,13 @@ class ServerEventsCog(commands.Cog):
|
|||||||
|
|
||||||
@commands.Cog.listener()
|
@commands.Cog.listener()
|
||||||
async def on_guild_role_update(
|
async def on_guild_role_update(
|
||||||
self, before: discord.Role, after: discord.Role
|
self,
|
||||||
|
before: discord.Role,
|
||||||
|
after: discord.Role,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Queue all text channels for ingestion when role permissions change.
|
||||||
Queues all text channels for message history ingestion when role permissions change,
|
|
||||||
but only if the bot is a member of the updated role.
|
Only triggers if the bot is a member of the updated role.
|
||||||
|
|
||||||
:param before: The role before the update.
|
:param before: The role before the update.
|
||||||
:param after: The role after the update.
|
:param after: The role after the update.
|
||||||
@@ -105,11 +104,11 @@ class ServerEventsCog(commands.Cog):
|
|||||||
|
|
||||||
@commands.Cog.listener()
|
@commands.Cog.listener()
|
||||||
async def on_guild_channel_update(
|
async def on_guild_channel_update(
|
||||||
self, before: discord.abc.GuildChannel, after: discord.abc.GuildChannel
|
self,
|
||||||
|
before: discord.abc.GuildChannel,
|
||||||
|
after: discord.abc.GuildChannel,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Queue all text channels for ingestion on permission change.
|
||||||
Queues all text channels for message history ingestion when channel override permissions
|
|
||||||
change.
|
|
||||||
|
|
||||||
:param before: The channel before the update.
|
:param before: The channel before the update.
|
||||||
:param after: The channel after the update.
|
:param after: The channel after the update.
|
||||||
@@ -121,8 +120,7 @@ class ServerEventsCog(commands.Cog):
|
|||||||
|
|
||||||
|
|
||||||
async def setup(bot: Crabstero) -> None:
|
async def setup(bot: Crabstero) -> None:
|
||||||
"""
|
"""Add the ServerEventsCog to the bot.
|
||||||
Adds the ServerEventsCog to the bot.
|
|
||||||
|
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
"""
|
"""
|
||||||
|
|||||||
+127
-63
@@ -12,16 +12,20 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Markov chain sentence ingestion and generation backed by SQLite.
|
||||||
Ingests sentences and generates new ones using a Markov chain backed by SQLite storage.
|
|
||||||
|
|
||||||
The Markov chain stores word transitions per channel, using duplicate rows to represent frequency
|
The Markov chain stores word transitions per channel, using duplicate rows
|
||||||
weight. Random selection via ORDER BY RANDOM() LIMIT 1 naturally preserves this weighting.
|
to represent frequency weight. Random selection via ORDER BY RANDOM()
|
||||||
|
LIMIT 1 naturally preserves this weighting.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import itertools
|
||||||
import re
|
import re
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from crabstero import metrics
|
||||||
|
from crabstero.database import StartWord, Transition
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
|
||||||
@@ -29,11 +33,24 @@ if TYPE_CHECKING:
|
|||||||
# as is common on Discord.
|
# as is common on Discord.
|
||||||
DEFAULT_SENTENCE_END = "\u00a7"
|
DEFAULT_SENTENCE_END = "\u00a7"
|
||||||
|
|
||||||
|
_TERMINATORS = frozenset({DEFAULT_SENTENCE_END, ".", "!", "?"})
|
||||||
|
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?]) ")
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize_whitespace(text: str) -> str:
|
||||||
|
"""Collapse runs of whitespace into single spaces and strip edges.
|
||||||
|
|
||||||
|
:param text: The raw text to normalize.
|
||||||
|
:return: The normalized text.
|
||||||
|
"""
|
||||||
|
return " ".join(text.split())
|
||||||
|
|
||||||
|
|
||||||
def is_complete_sentence(sentence: str) -> bool:
|
def is_complete_sentence(sentence: str) -> bool:
|
||||||
"""
|
"""Check whether a sentence ends with a valid terminator.
|
||||||
Checks whether a given sentence ends with a default sentence end character, a period, an
|
|
||||||
exclamation mark, or a question mark.
|
Valid terminators are the section sign, period, exclamation mark, or
|
||||||
|
question mark.
|
||||||
|
|
||||||
:param sentence: The sentence to test.
|
:param sentence: The sentence to test.
|
||||||
:return: True if the sentence ends with a valid terminator, False otherwise.
|
:return: True if the sentence ends with a valid terminator, False otherwise.
|
||||||
@@ -41,84 +58,136 @@ def is_complete_sentence(sentence: str) -> bool:
|
|||||||
if not sentence:
|
if not sentence:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
return sentence[-1] in (DEFAULT_SENTENCE_END, ".", "!", "?")
|
return sentence[-1] in _TERMINATORS
|
||||||
|
|
||||||
|
|
||||||
async def ingest(db: Database, channel_id: int, paragraph: str) -> None:
|
def _split_sentences(paragraph: str) -> list[str]:
|
||||||
"""
|
"""Normalize whitespace and split a paragraph into sentences.
|
||||||
Ingests a string potentially containing multiple smaller sentences into the Markov chain
|
|
||||||
for a given channel.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param paragraph: The raw paragraph text.
|
||||||
:param channel_id: The Discord channel ID to associate with this data.
|
:return: A list of individual sentences.
|
||||||
:param paragraph: The paragraph of sentences to ingest.
|
|
||||||
"""
|
"""
|
||||||
if not is_complete_sentence(paragraph):
|
if not is_complete_sentence(paragraph):
|
||||||
paragraph += DEFAULT_SENTENCE_END
|
paragraph += DEFAULT_SENTENCE_END
|
||||||
|
normalized = _normalize_whitespace(paragraph)
|
||||||
# Normalize whitespace, then split on sentence-ending punctuation followed by a space.
|
return _SENTENCE_SPLIT.split(normalized)
|
||||||
normalized = re.sub(r" +", " ", paragraph.strip().replace("\n", " "))
|
|
||||||
sentences = re.split(r"(?<=[.!?]) ", normalized)
|
|
||||||
|
|
||||||
for sentence in sentences:
|
|
||||||
await _ingest_sentence(db, channel_id, sentence)
|
|
||||||
|
|
||||||
|
|
||||||
async def _ingest_sentence(db: Database, channel_id: int, sentence: str) -> None:
|
_MIN_WORDS_FOR_START = 2
|
||||||
"""
|
|
||||||
Ingests a string containing a single sentence into the Markov chain for a given channel.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
|
||||||
:param channel_id: The Discord channel ID to associate with this data.
|
def _tokenize_sentence(
|
||||||
:param sentence: The sentence to ingest.
|
sentence: str,
|
||||||
|
) -> tuple[list[str], list[tuple[str, str]]]:
|
||||||
|
"""Tokenize a sentence into start words and transition pairs.
|
||||||
|
|
||||||
|
:param sentence: A single sentence.
|
||||||
|
:return: A tuple of (start_words, transitions) using bare word strings.
|
||||||
"""
|
"""
|
||||||
if not is_complete_sentence(sentence):
|
if not is_complete_sentence(sentence):
|
||||||
sentence += DEFAULT_SENTENCE_END
|
sentence += DEFAULT_SENTENCE_END
|
||||||
|
words = sentence.split()
|
||||||
|
|
||||||
# Normalize whitespace and split into individual words.
|
start_words: list[str] = [words[0]] if len(words) >= _MIN_WORDS_FOR_START else []
|
||||||
words = re.sub(r" +", " ", sentence.strip()).split(" ")
|
transitions: list[tuple[str, str]] = list(itertools.pairwise(words))
|
||||||
|
return start_words, transitions
|
||||||
|
|
||||||
start_words: list[tuple[int, str]] = []
|
|
||||||
transitions: list[tuple[int, str, str]] = []
|
|
||||||
|
|
||||||
for i in range(len(words) - 1):
|
async def ingest(db: Database, channel_id: int, user_id: int, paragraph: str) -> None:
|
||||||
if i == 0:
|
"""Ingest a paragraph into the Markov chain for a given channel.
|
||||||
start_words.append((channel_id, words[i]))
|
|
||||||
transitions.append((channel_id, words[i], words[i + 1]))
|
|
||||||
|
|
||||||
if start_words:
|
The paragraph may contain multiple sentences which are split and ingested
|
||||||
await db.add_start_words_batch(start_words)
|
individually.
|
||||||
if transitions:
|
|
||||||
await db.add_transitions_batch(transitions)
|
:param db: The database instance.
|
||||||
|
:param channel_id: The Discord channel ID to associate with this data.
|
||||||
|
:param user_id: The Discord user ID of the contributor.
|
||||||
|
:param paragraph: The paragraph of sentences to ingest.
|
||||||
|
"""
|
||||||
|
for sentence in _split_sentences(paragraph):
|
||||||
|
await _ingest_sentence(db, channel_id, user_id, sentence)
|
||||||
|
|
||||||
|
|
||||||
|
async def _ingest_sentence(
|
||||||
|
db: Database,
|
||||||
|
channel_id: int,
|
||||||
|
user_id: int,
|
||||||
|
sentence: str,
|
||||||
|
) -> None:
|
||||||
|
"""Ingest a single sentence into the Markov chain for a given channel.
|
||||||
|
|
||||||
|
:param db: The database instance.
|
||||||
|
:param channel_id: The Discord channel ID to associate with this data.
|
||||||
|
:param user_id: The Discord user ID of the contributor.
|
||||||
|
:param sentence: The sentence to ingest.
|
||||||
|
"""
|
||||||
|
raw_starts, raw_transitions = _tokenize_sentence(sentence)
|
||||||
|
start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
|
||||||
|
transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
|
||||||
|
await db.add_markov_data(start_words, transitions)
|
||||||
|
|
||||||
|
|
||||||
|
async def uningest(db: Database, channel_id: int, user_id: int, paragraph: str) -> None:
|
||||||
|
"""Remove a paragraph's Markov data from the chain for a given channel.
|
||||||
|
|
||||||
|
Mirrors :func:`ingest` but deletes one matching row per entry instead of
|
||||||
|
inserting.
|
||||||
|
|
||||||
|
:param db: The database instance.
|
||||||
|
:param channel_id: The Discord channel ID.
|
||||||
|
:param user_id: The Discord user ID of the contributor.
|
||||||
|
:param paragraph: The paragraph of sentences to uningest.
|
||||||
|
"""
|
||||||
|
for sentence in _split_sentences(paragraph):
|
||||||
|
await _uningest_sentence(db, channel_id, user_id, sentence)
|
||||||
|
|
||||||
|
|
||||||
|
async def _uningest_sentence(
|
||||||
|
db: Database,
|
||||||
|
channel_id: int,
|
||||||
|
user_id: int,
|
||||||
|
sentence: str,
|
||||||
|
) -> None:
|
||||||
|
"""Remove a single sentence's Markov data from the chain.
|
||||||
|
|
||||||
|
:param db: The database instance.
|
||||||
|
:param channel_id: The Discord channel ID.
|
||||||
|
:param user_id: The Discord user ID of the contributor.
|
||||||
|
:param sentence: The sentence to uningest.
|
||||||
|
"""
|
||||||
|
raw_starts, raw_transitions = _tokenize_sentence(sentence)
|
||||||
|
start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
|
||||||
|
transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
|
||||||
|
await db.remove_markov_data(start_words, transitions)
|
||||||
|
|
||||||
|
|
||||||
async def generate(
|
async def generate(
|
||||||
db: Database, channel_id: int, soft_limit: int = 750, hard_limit: int = 1000
|
db: Database,
|
||||||
|
channel_id: int,
|
||||||
|
soft_limit: int = 750,
|
||||||
|
hard_limit: int = 1000,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""Generate a new sentence from previously ingested words for a channel.
|
||||||
Generates a new sentence using words learned from previously ingested sentences for a given
|
|
||||||
channel.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param channel_id: The Discord channel ID to generate from.
|
:param channel_id: The Discord channel ID to generate from.
|
||||||
:param soft_limit: The amount of characters to try and limit sentence length around.
|
:param soft_limit: Character count to aim for when wrapping up.
|
||||||
:param hard_limit: The amount of characters to cut off the sentence at if it gets too long.
|
:param hard_limit: Character count to hard-cut the sentence at.
|
||||||
:return: A new generated sentence.
|
:return: A new generated sentence.
|
||||||
"""
|
"""
|
||||||
word = await db.get_random_start_word(channel_id)
|
word = await db.get_random_start_word(channel_id)
|
||||||
|
|
||||||
# Seed the chain with a fallback sentence if the channel has no data yet.
|
|
||||||
if word is None:
|
if word is None:
|
||||||
await _ingest_sentence(db, channel_id, "Hello world!")
|
return (
|
||||||
word = await db.get_random_start_word(channel_id)
|
"I do not have enough data to generate a message yet."
|
||||||
if word is None:
|
" Chat a bit more so I can learn how this channel talks."
|
||||||
return ""
|
)
|
||||||
|
|
||||||
parts: list[str] = []
|
with metrics.GENERATION_DURATION.time():
|
||||||
parts.append(word)
|
parts: list[str] = [word]
|
||||||
current_length = len(word)
|
current_length = len(word)
|
||||||
|
|
||||||
# The loop is skipped if the starting word already ends a sentence (e.g. "Yes.").
|
# The loop is skipped if the start word already ends a sentence (e.g. "Yes.").
|
||||||
while not is_complete_sentence(word):
|
while not is_complete_sentence(word):
|
||||||
# Past the soft limit, prefer a sentence-ending word to wrap up.
|
# Past the soft limit, prefer a sentence-ending word to wrap up.
|
||||||
if current_length >= soft_limit:
|
if current_length >= soft_limit:
|
||||||
@@ -137,16 +206,11 @@ async def generate(
|
|||||||
current_length += 1 + len(word) # +1 for the joining space.
|
current_length += 1 + len(word) # +1 for the joining space.
|
||||||
|
|
||||||
if current_length >= hard_limit:
|
if current_length >= hard_limit:
|
||||||
result = " ".join(parts)[:hard_limit]
|
break
|
||||||
# Strip the internal sentence-end marker if it ended up at the boundary.
|
|
||||||
if result and result[-1] == DEFAULT_SENTENCE_END:
|
|
||||||
return result[:-1]
|
|
||||||
return result
|
|
||||||
|
|
||||||
result = " ".join(parts)
|
result = " ".join(parts)
|
||||||
|
if current_length >= hard_limit:
|
||||||
|
result = result[:hard_limit]
|
||||||
|
|
||||||
# Strip the internal sentence-end marker so it never appears in output.
|
# Strip the internal sentence-end marker so it never appears in output.
|
||||||
if result and result[-1] == DEFAULT_SENTENCE_END:
|
return result.removesuffix(DEFAULT_SENTENCE_END)
|
||||||
return result[:-1]
|
|
||||||
|
|
||||||
return result
|
|
||||||
|
|||||||
+104
-33
@@ -12,11 +12,11 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Generate Discord reply messages and ingest raw messages.
|
||||||
Assists with generating Discord messages in response to other users and ingesting raw messages.
|
|
||||||
|
|
||||||
Orchestrates Markov chain generation and ingestion in the context of Discord messages, handling
|
Orchestrates Markov chain generation and ingestion in the context of
|
||||||
reply logic, embed generation, mention filtering, and message ingestion with flag checks.
|
Discord messages, handling reply logic, embed generation, mention
|
||||||
|
filtering, and message ingestion with flag checks.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import re
|
import re
|
||||||
@@ -25,22 +25,23 @@ from typing import TYPE_CHECKING
|
|||||||
|
|
||||||
import discord
|
import discord
|
||||||
|
|
||||||
from crabstero import flags, markov
|
from crabstero import flags, markov, metrics
|
||||||
|
from crabstero.cache import CachedMessage, IngestCache
|
||||||
|
from crabstero.database import ChannelImage
|
||||||
from crabstero.flags import EntityType, Flag
|
from crabstero.flags import EntityType, Flag
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
|
||||||
# Copied from discordjs/discord-api-types:
|
# Copied from discordjs/discord-api-types:
|
||||||
# https://github.com/discordjs/discord-api-types/blob/7fe434114e91c80ed79f0204ae6c73047672d55d/globals.ts#L30
|
# https://github.com/discordjs/discord-api-types/blob/662cb0cb0ac9c6f9ad93e180849476714bfceb0c/globals.ts#L39
|
||||||
MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
|
_MENTION_PATTERN = re.compile(r"<@!?(\d{17,20})>")
|
||||||
|
|
||||||
_EMBED_CHANCE_THRESHOLD = 95 # Out of 100; sends an embed ~5% of the time.
|
_EMBED_CHANCE_THRESHOLD = 95 # Out of 100; sends an embed ~5% of the time.
|
||||||
|
|
||||||
|
|
||||||
async def reply_to_message(db: Database, message: discord.Message) -> None:
|
async def reply_to_message(db: Database, message: discord.Message) -> None:
|
||||||
"""
|
"""Send a new message in Discord in response to a given message.
|
||||||
Sends a new message in Discord in response to a given message.
|
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param message: The message prompting the response.
|
:param message: The message prompting the response.
|
||||||
@@ -66,7 +67,7 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
|
|||||||
else:
|
else:
|
||||||
channel_id = channel.id
|
channel_id = channel.id
|
||||||
|
|
||||||
body = await markov.generate(db, channel_id, 750, 1000)
|
body = await markov.generate(db, channel_id)
|
||||||
|
|
||||||
embed = None
|
embed = None
|
||||||
|
|
||||||
@@ -76,8 +77,13 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
|
|||||||
and channel.permissions_for(guild.me).embed_links
|
and channel.permissions_for(guild.me).embed_links
|
||||||
):
|
):
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
title=await markov.generate(db, channel_id, 200, 300),
|
title=await markov.generate(db, channel_id, soft_limit=200, hard_limit=300),
|
||||||
description=await markov.generate(db, channel_id, 300, 500),
|
description=await markov.generate(
|
||||||
|
db,
|
||||||
|
channel_id,
|
||||||
|
soft_limit=300,
|
||||||
|
hard_limit=500,
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
random_image = await db.get_random_image(channel_id)
|
random_image = await db.get_random_image(channel_id)
|
||||||
@@ -85,13 +91,11 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
|
|||||||
embed.set_image(url=random_image)
|
embed.set_image(url=random_image)
|
||||||
|
|
||||||
# Suppress all mentions by default; only ping users who opted in.
|
# Suppress all mentions by default; only ping users who opted in.
|
||||||
allowed_user_ids: list[int] = []
|
allowed_user_ids: set[int] = set()
|
||||||
|
|
||||||
for match in MENTION_PATTERN.finditer(body):
|
|
||||||
user_id = int(match.group("id"))
|
|
||||||
|
|
||||||
|
for user_id in {int(uid) for uid in _MENTION_PATTERN.findall(body)}:
|
||||||
if await flags.is_flag_set(db, user_id, EntityType.USER, Flag.ALLOW_PINGS):
|
if await flags.is_flag_set(db, user_id, EntityType.USER, Flag.ALLOW_PINGS):
|
||||||
allowed_user_ids.append(user_id)
|
allowed_user_ids.add(user_id)
|
||||||
|
|
||||||
allowed_mentions = discord.AllowedMentions(
|
allowed_mentions = discord.AllowedMentions(
|
||||||
everyone=False,
|
everyone=False,
|
||||||
@@ -106,20 +110,50 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
|
|||||||
allowed_mentions=allowed_mentions,
|
allowed_mentions=allowed_mentions,
|
||||||
mention_author=False,
|
mention_author=False,
|
||||||
)
|
)
|
||||||
|
metrics.EMBEDS_GENERATED.inc()
|
||||||
else:
|
else:
|
||||||
await message.reply(
|
await message.reply(
|
||||||
content=body,
|
content=body,
|
||||||
allowed_mentions=allowed_mentions,
|
allowed_mentions=allowed_mentions,
|
||||||
mention_author=False,
|
mention_author=False,
|
||||||
)
|
)
|
||||||
|
metrics.REPLIES_SENT.inc()
|
||||||
|
|
||||||
|
|
||||||
async def ingest_message(db: Database, message: discord.Message) -> None:
|
def _extract_embed_data(
|
||||||
|
embeds: list[discord.Embed],
|
||||||
|
) -> tuple[list[str], list[str]]:
|
||||||
|
"""Extract text and image URLs from message embeds.
|
||||||
|
|
||||||
|
:param embeds: The embeds to process.
|
||||||
|
:return: A tuple of (texts, image_urls).
|
||||||
"""
|
"""
|
||||||
Ingests a given message into its channel's Markov chain.
|
texts: list[str] = []
|
||||||
|
image_urls: list[str] = []
|
||||||
|
for embed in embeds:
|
||||||
|
if embed.title:
|
||||||
|
texts.append(embed.title)
|
||||||
|
if embed.description:
|
||||||
|
texts.append(embed.description)
|
||||||
|
if embed.image and embed.image.url:
|
||||||
|
image_urls.append(embed.image.url)
|
||||||
|
return texts, image_urls
|
||||||
|
|
||||||
|
|
||||||
|
async def ingest_message(
|
||||||
|
db: Database,
|
||||||
|
message: discord.Message,
|
||||||
|
cache: IngestCache | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Ingest a given message into its channel's Markov chain.
|
||||||
|
|
||||||
|
If a cache is provided, the ingested data is recorded so it can be
|
||||||
|
reversed by :func:`uningest_message` if the message is deleted shortly
|
||||||
|
after.
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param message: The message to ingest.
|
:param message: The message to ingest.
|
||||||
|
:param cache: Optional ingest cache for tracking recently ingested data.
|
||||||
"""
|
"""
|
||||||
guild = message.guild
|
guild = message.guild
|
||||||
if guild is None:
|
if guild is None:
|
||||||
@@ -136,27 +170,64 @@ async def ingest_message(db: Database, message: discord.Message) -> None:
|
|||||||
return
|
return
|
||||||
|
|
||||||
channel_id = message.channel.id
|
channel_id = message.channel.id
|
||||||
|
user_id = message.author.id
|
||||||
|
|
||||||
|
if not message.content and not message.embeds:
|
||||||
|
return
|
||||||
|
|
||||||
|
embed_texts, image_urls = _extract_embed_data(message.embeds)
|
||||||
|
|
||||||
|
with metrics.MESSAGE_INGESTION_DURATION.time():
|
||||||
if message.content:
|
if message.content:
|
||||||
await markov.ingest(db, channel_id, message.content)
|
await markov.ingest(db, channel_id, user_id, message.content)
|
||||||
|
|
||||||
for embed in message.embeds:
|
for text in embed_texts:
|
||||||
await _ingest_embed(db, channel_id, embed)
|
await markov.ingest(db, channel_id, user_id, text)
|
||||||
|
|
||||||
|
if image_urls:
|
||||||
|
await db.add_images(
|
||||||
|
[ChannelImage(channel_id, user_id, url) for url in image_urls],
|
||||||
|
)
|
||||||
|
|
||||||
|
if cache is not None:
|
||||||
|
cache.put(
|
||||||
|
message.id,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=channel_id,
|
||||||
|
user_id=user_id,
|
||||||
|
content=message.content or None,
|
||||||
|
embed_texts=embed_texts,
|
||||||
|
image_urls=image_urls,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def _ingest_embed(db: Database, channel_id: int, embed: discord.Embed) -> None:
|
async def uningest_message(db: Database, cache: IngestCache, message_id: int) -> None:
|
||||||
"""
|
"""Reverse ingestion for a recently deleted message.
|
||||||
Ingests a given embed into a given channel's Markov chain.
|
|
||||||
|
Looks up the message in the cache. If found, removes all Markov data
|
||||||
|
and images that were added during ingestion.
|
||||||
|
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
:param channel_id: The ID of the channel to use for the Markov chain.
|
:param cache: The ingest cache.
|
||||||
:param embed: The embed to ingest.
|
:param message_id: The Discord message ID that was deleted.
|
||||||
"""
|
"""
|
||||||
if embed.title:
|
entry = cache.pop(message_id)
|
||||||
await markov.ingest(db, channel_id, embed.title)
|
if entry is None:
|
||||||
|
return
|
||||||
|
|
||||||
if embed.description:
|
if entry.content:
|
||||||
await markov.ingest(db, channel_id, embed.description)
|
await markov.uningest(db, entry.channel_id, entry.user_id, entry.content)
|
||||||
|
|
||||||
if embed.image and embed.image.url:
|
for text in entry.embed_texts:
|
||||||
await db.add_image(channel_id, embed.image.url)
|
await markov.uningest(db, entry.channel_id, entry.user_id, text)
|
||||||
|
|
||||||
|
if entry.image_urls:
|
||||||
|
await db.remove_images(
|
||||||
|
[
|
||||||
|
ChannelImage(entry.channel_id, entry.user_id, url)
|
||||||
|
for url in entry.image_urls
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
metrics.MESSAGES_UNINGESTED.inc()
|
||||||
|
|||||||
@@ -0,0 +1,254 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Prometheus metrics definitions and HTTP server for Crabstero.
|
||||||
|
|
||||||
|
All metric objects are defined at module level using the default global
|
||||||
|
registry. The MetricsServer class wraps an aiohttp application that serves
|
||||||
|
the ``/metrics`` scrape endpoint.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import contextlib
|
||||||
|
import errno
|
||||||
|
import logging
|
||||||
|
import socket
|
||||||
|
import stat
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from aiohttp import web
|
||||||
|
from prometheus_client import Counter, Gauge, Histogram, Info
|
||||||
|
from prometheus_client.aiohttp import make_aiohttp_handler
|
||||||
|
|
||||||
|
from crabstero import __version__
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_FAST_BUCKETS = (0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0)
|
||||||
|
_SLOW_BUCKETS = (0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0)
|
||||||
|
_MESSAGES_BUCKETS = (100, 500, 1000, 5000, 10000, 25000, 50000)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TcpMetricsAddress:
|
||||||
|
"""TCP listen address for the metrics server."""
|
||||||
|
|
||||||
|
host: str
|
||||||
|
port: int
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class UnixMetricsAddress:
|
||||||
|
"""Unix socket listen address for the metrics server."""
|
||||||
|
|
||||||
|
path: str
|
||||||
|
|
||||||
|
|
||||||
|
type MetricsAddress = TcpMetricsAddress | UnixMetricsAddress
|
||||||
|
|
||||||
|
|
||||||
|
BUILD_INFO = Info("crabstero_build", "Build information")
|
||||||
|
BUILD_INFO.info({"version": __version__})
|
||||||
|
|
||||||
|
MESSAGES_UNINGESTED = Counter(
|
||||||
|
"crabstero_messages_uningested_total",
|
||||||
|
"Messages uningested from a Markov chain after deletion",
|
||||||
|
)
|
||||||
|
REPLIES_SENT = Counter(
|
||||||
|
"crabstero_replies_sent_total",
|
||||||
|
"Markov replies sent",
|
||||||
|
)
|
||||||
|
EMBEDS_GENERATED = Counter(
|
||||||
|
"crabstero_embeds_generated_total",
|
||||||
|
"Embeds included in replies",
|
||||||
|
)
|
||||||
|
MESSAGES_PROCESSED = Counter(
|
||||||
|
"crabstero_messages_processed_total",
|
||||||
|
"Messages seen by the listener",
|
||||||
|
)
|
||||||
|
PINGME = Counter(
|
||||||
|
"crabstero_pingme_total",
|
||||||
|
"/pingme command outcomes",
|
||||||
|
["outcome"],
|
||||||
|
)
|
||||||
|
PINGME.labels(outcome="opted_in")
|
||||||
|
PINGME.labels(outcome="opted_out")
|
||||||
|
|
||||||
|
FORGETME = Counter(
|
||||||
|
"crabstero_forgetme_total",
|
||||||
|
"/forgetme command outcomes",
|
||||||
|
["outcome"],
|
||||||
|
)
|
||||||
|
FORGETME.labels(outcome="completed")
|
||||||
|
FORGETME.labels(outcome="cancelled")
|
||||||
|
FORGETME.labels(outcome="timeout")
|
||||||
|
FORGETME.labels(outcome="already_forgotten")
|
||||||
|
|
||||||
|
ERRORS = Counter(
|
||||||
|
"crabstero_errors_total",
|
||||||
|
"Unhandled exceptions by source",
|
||||||
|
["source"],
|
||||||
|
)
|
||||||
|
DISCORD_EVENTS = Counter(
|
||||||
|
"crabstero_discord_events_total",
|
||||||
|
"Discord gateway events received",
|
||||||
|
["event"],
|
||||||
|
)
|
||||||
|
|
||||||
|
DISCORD_LATENCY = Gauge(
|
||||||
|
"crabstero_discord_latency_seconds",
|
||||||
|
"Discord WebSocket heartbeat latency",
|
||||||
|
)
|
||||||
|
GUILD_COUNT = Gauge(
|
||||||
|
"crabstero_guild_count",
|
||||||
|
"Guilds the bot is currently in",
|
||||||
|
)
|
||||||
|
INGESTION_ACTIVE = Gauge(
|
||||||
|
"crabstero_ingestion_active_channels",
|
||||||
|
"Channels pending or in-progress for ingestion",
|
||||||
|
)
|
||||||
|
|
||||||
|
MESSAGE_INGESTION_DURATION = Histogram(
|
||||||
|
"crabstero_message_ingestion_duration_seconds",
|
||||||
|
"Time to ingest one message",
|
||||||
|
buckets=_FAST_BUCKETS,
|
||||||
|
)
|
||||||
|
CHANNEL_INGESTION_DURATION = Histogram(
|
||||||
|
"crabstero_channel_ingestion_duration_seconds",
|
||||||
|
"Time to fully ingest a channel",
|
||||||
|
buckets=_SLOW_BUCKETS,
|
||||||
|
)
|
||||||
|
GENERATION_DURATION = Histogram(
|
||||||
|
"crabstero_generation_duration_seconds",
|
||||||
|
"Time to generate a Markov response",
|
||||||
|
buckets=_FAST_BUCKETS,
|
||||||
|
)
|
||||||
|
CHANNEL_INGESTION_MESSAGES = Histogram(
|
||||||
|
"crabstero_channel_ingestion_messages",
|
||||||
|
"Messages iterated per channel ingestion run",
|
||||||
|
buckets=_MESSAGES_BUCKETS,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MetricsServer:
|
||||||
|
"""HTTP server that exposes a Prometheus ``/metrics`` scrape endpoint.
|
||||||
|
|
||||||
|
Uses ``aiohttp.web.AppRunner`` with TCP or Unix socket sites for
|
||||||
|
async-native serving.
|
||||||
|
The handler is provided by ``prometheus_client.aiohttp.make_aiohttp_handler``,
|
||||||
|
which handles compression and content negotiation automatically.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__slots__ = ("_address", "_port", "_runner")
|
||||||
|
|
||||||
|
def __init__(self, address: MetricsAddress) -> None:
|
||||||
|
"""Store the listen address.
|
||||||
|
|
||||||
|
:param address: TCP or Unix socket address to bind to.
|
||||||
|
"""
|
||||||
|
self._address = address
|
||||||
|
self._port = address.port if isinstance(address, TcpMetricsAddress) else None
|
||||||
|
self._runner: web.AppRunner | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def port(self) -> int:
|
||||||
|
"""The TCP port the server is bound to.
|
||||||
|
|
||||||
|
After ``start()``, this reflects the actual port (useful when
|
||||||
|
the constructor received port 0 for OS assignment).
|
||||||
|
"""
|
||||||
|
if self._port is None:
|
||||||
|
msg = "Unix socket metrics server does not have a TCP port"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
return self._port
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
"""Create the aiohttp application and start listening."""
|
||||||
|
if self._runner is not None:
|
||||||
|
msg = "Metrics server is already running"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
app = web.Application()
|
||||||
|
app.router.add_get("/metrics", make_aiohttp_handler())
|
||||||
|
self._runner = web.AppRunner(app)
|
||||||
|
await self._runner.setup()
|
||||||
|
try:
|
||||||
|
if isinstance(self._address, TcpMetricsAddress):
|
||||||
|
site: web.BaseSite = web.TCPSite(
|
||||||
|
self._runner,
|
||||||
|
self._address.host,
|
||||||
|
self._address.port,
|
||||||
|
)
|
||||||
|
await site.start()
|
||||||
|
|
||||||
|
# Resolve the actual bound port when the OS assigned one.
|
||||||
|
if self._address.port == 0:
|
||||||
|
self._port = self._runner.addresses[0][1]
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Metrics server listening on %s:%d.",
|
||||||
|
self._address.host,
|
||||||
|
self.port,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
socket_path = Path(self._address.path)
|
||||||
|
# This runs once before binding the Unix site, so synchronous
|
||||||
|
# path checks are simpler than offloading startup setup.
|
||||||
|
try:
|
||||||
|
mode = socket_path.stat( # noqa: ASYNC240
|
||||||
|
follow_symlinks=False,
|
||||||
|
).st_mode
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
if not stat.S_ISSOCK(mode):
|
||||||
|
msg = (
|
||||||
|
"Metrics Unix socket path already exists and is not a "
|
||||||
|
f"socket: {self._address.path}"
|
||||||
|
)
|
||||||
|
raise FileExistsError(msg)
|
||||||
|
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
|
||||||
|
probe.settimeout(0.1)
|
||||||
|
try:
|
||||||
|
probe.connect(self._address.path)
|
||||||
|
except OSError as exc:
|
||||||
|
if exc.errno not in {errno.ECONNREFUSED, errno.ENOENT}:
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
msg = (
|
||||||
|
"Metrics Unix socket path is already in use: "
|
||||||
|
f"{self._address.path}"
|
||||||
|
)
|
||||||
|
raise OSError(errno.EADDRINUSE, msg)
|
||||||
|
|
||||||
|
with contextlib.suppress(FileNotFoundError):
|
||||||
|
socket_path.unlink() # noqa: ASYNC240
|
||||||
|
|
||||||
|
site = web.UnixSite(self._runner, self._address.path)
|
||||||
|
await site.start()
|
||||||
|
logger.info(
|
||||||
|
"Metrics server listening on Unix socket %s.",
|
||||||
|
self._address.path,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
await self._runner.cleanup()
|
||||||
|
self._runner = None
|
||||||
|
raise
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""Shut down the HTTP server and release resources."""
|
||||||
|
if self._runner is not None:
|
||||||
|
await self._runner.cleanup()
|
||||||
|
self._runner = None
|
||||||
|
logger.info("Metrics server stopped.")
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Background tasks for Crabstero."""
|
||||||
|
|||||||
@@ -12,10 +12,9 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
"""
|
"""Bulk-ingests the message history of channels.
|
||||||
Bulk-ingests the message history of channels.
|
|
||||||
|
|
||||||
Channels are enqueued for processing by the bot's fixed ingestion worker pool.
|
Channels are processed by the bot's semaphore-bounded dynamic tasks.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
@@ -23,10 +22,11 @@ from typing import TYPE_CHECKING
|
|||||||
|
|
||||||
import discord
|
import discord
|
||||||
|
|
||||||
|
from crabstero import metrics
|
||||||
from crabstero.messages import ingest_message
|
from crabstero.messages import ingest_message
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crabstero.bot import Crabstero
|
from crabstero.bot import Crabstero, IngestableChannel
|
||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
|
||||||
MAXIMUM_MESSAGES_PER_CHANNEL = (
|
MAXIMUM_MESSAGES_PER_CHANNEL = (
|
||||||
@@ -37,8 +37,7 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
|
def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
|
||||||
"""
|
"""Enqueue every textable channel in a guild for message history ingestion.
|
||||||
Enqueues every textable channel in a guild for message history ingestion.
|
|
||||||
|
|
||||||
:param guild: The Discord guild whose channels should be ingested.
|
:param guild: The Discord guild whose channels should be ingested.
|
||||||
:param bot: The bot instance.
|
:param bot: The bot instance.
|
||||||
@@ -49,20 +48,21 @@ def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
|
|||||||
|
|
||||||
|
|
||||||
async def ingest_channel(
|
async def ingest_channel(
|
||||||
channel: discord.TextChannel | discord.VoiceChannel,
|
channel: IngestableChannel,
|
||||||
db: Database,
|
db: Database,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""Bulk-ingest the message history of a given channel.
|
||||||
Bulk-ingests the message history of a given channel. The task will be ended early if
|
|
||||||
permissions do not allow ingesting this channel or if it has already been ingested in the past.
|
End early if permissions do not allow ingesting this channel or if it has
|
||||||
|
already been ingested.
|
||||||
|
|
||||||
:param channel: The channel to ingest.
|
:param channel: The channel to ingest.
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
"""
|
"""
|
||||||
try:
|
|
||||||
if not channel.permissions_for(channel.guild.me).read_message_history:
|
if not channel.permissions_for(channel.guild.me).read_message_history:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"[%s] Unable to ingest textable channel history due to lacking permissions. Ignoring.",
|
"[%s] Unable to ingest channel history"
|
||||||
|
" due to lacking permissions. Ignoring.",
|
||||||
channel.id,
|
channel.id,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
@@ -74,19 +74,16 @@ async def ingest_channel(
|
|||||||
|
|
||||||
logger.info("[%s] Starting ingestion of textable channel history.", channel.id)
|
logger.info("[%s] Starting ingestion of textable channel history.", channel.id)
|
||||||
|
|
||||||
|
with metrics.CHANNEL_INGESTION_DURATION.time():
|
||||||
count = 0
|
count = 0
|
||||||
async for message in channel.history(limit=MAXIMUM_MESSAGES_PER_CHANNEL):
|
async for message in channel.history(limit=MAXIMUM_MESSAGES_PER_CHANNEL):
|
||||||
count += 1
|
count += 1
|
||||||
await ingest_message(db, message)
|
await ingest_message(db, message)
|
||||||
|
|
||||||
|
metrics.CHANNEL_INGESTION_MESSAGES.observe(count)
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"[%s] Ingestion of textable channel history complete. %d messages ingested.",
|
"[%s] Ingestion of channel history complete. %d messages iterated.",
|
||||||
channel.id,
|
channel.id,
|
||||||
count,
|
count,
|
||||||
)
|
)
|
||||||
|
|
||||||
except Exception:
|
|
||||||
logger.exception(
|
|
||||||
"[%s] An error occurred while ingesting textable channel history!",
|
|
||||||
channel.id,
|
|
||||||
)
|
|
||||||
|
|||||||
+41
-22
@@ -2,14 +2,17 @@
|
|||||||
name = "crabstero"
|
name = "crabstero"
|
||||||
dynamic = ["version"]
|
dynamic = ["version"]
|
||||||
description = "The simple nonversation Discord bot."
|
description = "The simple nonversation Discord bot."
|
||||||
license = "Apache-2.0"
|
|
||||||
requires-python = ">=3.14"
|
requires-python = ">=3.14"
|
||||||
|
license = "Apache-2.0"
|
||||||
authors = [
|
authors = [
|
||||||
{ name = "Logan Fick" },
|
{ name = "Logan Fick" },
|
||||||
]
|
]
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"discord.py>=2.6.4",
|
"discord.py>=2.7.1",
|
||||||
"aiosqlite>=0.22.1",
|
"aiosqlite>=0.22.1",
|
||||||
|
"aiohttp>=3.14.1",
|
||||||
|
"prometheus-client>=0.25.0",
|
||||||
|
"uvloop>=0.22.1",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
@@ -20,11 +23,13 @@ Repository = "https://git.logal.dev/LogalDeveloper/Crabstero"
|
|||||||
|
|
||||||
[dependency-groups]
|
[dependency-groups]
|
||||||
dev = [
|
dev = [
|
||||||
"pytest>=9.0.2",
|
"codespell>=2.4.2",
|
||||||
"pytest-asyncio>=1.3.0",
|
"mypy>=2.1.0",
|
||||||
"pytest-cov>=7.0.0",
|
"pip-audit>=2.10.1",
|
||||||
"mypy>=1.18.1",
|
"pytest>=9.0.3",
|
||||||
"ruff>=0.15.1",
|
"pytest-asyncio>=1.4.0",
|
||||||
|
"pytest-cov>=7.1.0",
|
||||||
|
"ruff>=0.15.16",
|
||||||
]
|
]
|
||||||
|
|
||||||
[tool.uv]
|
[tool.uv]
|
||||||
@@ -32,7 +37,7 @@ extra-index-url = ["https://git.logal.dev/api/packages/LogalDeveloper/pypi/simpl
|
|||||||
publish-url = "https://git.logal.dev/api/packages/LogalDeveloper/pypi"
|
publish-url = "https://git.logal.dev/api/packages/LogalDeveloper/pypi"
|
||||||
|
|
||||||
[build-system]
|
[build-system]
|
||||||
requires = ["hatchling>=1.28.0", "hatch-vcs>=0.5.0"]
|
requires = ["hatchling>=1.30.1", "hatch-vcs>=0.5.0"]
|
||||||
build-backend = "hatchling.build"
|
build-backend = "hatchling.build"
|
||||||
|
|
||||||
[tool.hatch.version]
|
[tool.hatch.version]
|
||||||
@@ -51,29 +56,43 @@ warn_unreachable = true
|
|||||||
explicit_package_bases = true
|
explicit_package_bases = true
|
||||||
|
|
||||||
[tool.pytest.ini_options]
|
[tool.pytest.ini_options]
|
||||||
testpaths = ["tests"]
|
|
||||||
asyncio_mode = "auto"
|
asyncio_mode = "auto"
|
||||||
|
asyncio_default_fixture_loop_scope = "function"
|
||||||
|
|
||||||
[tool.ruff]
|
[tool.ruff]
|
||||||
target-version = "py314"
|
target-version = "py314"
|
||||||
extend-exclude = ["crabstero/_version.py"]
|
extend-exclude = ["crabstero/_version.py"] # auto-generated by hatch-vcs
|
||||||
|
|
||||||
[tool.ruff.lint]
|
[tool.ruff.lint]
|
||||||
select = [
|
select = ["ALL"]
|
||||||
"F", # Pyflakes
|
|
||||||
"E", # pycodestyle errors
|
|
||||||
"W", # pycodestyle warnings
|
|
||||||
"I", # isort
|
|
||||||
"UP", # pyupgrade
|
|
||||||
"B", # flake8-bugbear
|
|
||||||
"SIM", # flake8-simplify
|
|
||||||
"TCH", # flake8-type-checking
|
|
||||||
"RUF", # Ruff-specific rules
|
|
||||||
]
|
|
||||||
ignore = [
|
ignore = [
|
||||||
"E501", # line length enforced by the formatter; residual violations are in strings/docstrings/comments
|
"ANN401", # Any is valid at system boundaries; mypy strict handles real issues
|
||||||
|
"C901", # McCabe complexity: noisy and not actionable
|
||||||
|
"COM812", # handled by the formatter
|
||||||
|
"D203", # incompatible with D211 (no blank line before class docstring)
|
||||||
|
"D213", # incompatible with D212 (summary on first line)
|
||||||
|
"EM", # exception message style: inline literals are fine
|
||||||
|
"PLR0911", # too many return statements: flat early-returns are clear
|
||||||
|
"PLR0912", # too many branches: inherent in parsers, validators, CLI
|
||||||
|
"PLR0913", # too many arguments: API surfaces and constructors need them
|
||||||
|
"PLR0915", # too many statements: inherent in parsers, validators, CLI
|
||||||
|
"TRY003", # inline exception messages are fine (complements EM ignore)
|
||||||
|
"TRY301", # raise inside try: guard clauses don't need helper functions
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[tool.ruff.lint.per-file-ignores]
|
||||||
|
"crabstero/__init__.py" = ["E402"] # version fallback is resolved before public imports
|
||||||
|
"tests/**" = [
|
||||||
|
"S101", # assert is standard for pytest
|
||||||
|
"S311", # pseudo-random generators are fine in tests
|
||||||
|
"SLF001", # tests legitimately access private members for verification
|
||||||
|
"ARG001", # unused args are normal for fixtures and handler stubs
|
||||||
|
"PLR2004", # magic values are clear in test assertions
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.codespell]
|
||||||
|
skip = "crabstero/_version.py,uv.lock"
|
||||||
|
|
||||||
[tool.coverage.run]
|
[tool.coverage.run]
|
||||||
source = ["crabstero"]
|
source = ["crabstero"]
|
||||||
omit = [
|
omit = [
|
||||||
|
|||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Test suite for Crabstero."""
|
||||||
|
|||||||
@@ -11,3 +11,22 @@
|
|||||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Shared pytest fixtures for the test suite."""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero.database import Database
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
async def db() -> AsyncGenerator[Database]:
|
||||||
|
"""Yield a Database backed by an in-memory SQLite database."""
|
||||||
|
database = await Database.connect(":memory:")
|
||||||
|
yield database
|
||||||
|
await database.close()
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for the IngestCache TTL cache."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time
|
||||||
|
|
||||||
|
from crabstero.cache import CachedMessage, IngestCache
|
||||||
|
|
||||||
|
|
||||||
|
class TestIngestCache:
|
||||||
|
"""IngestCache put, pop, and expiry behavior."""
|
||||||
|
|
||||||
|
def test_put_and_pop(self) -> None:
|
||||||
|
"""A cached message can be retrieved by message ID."""
|
||||||
|
cache = IngestCache()
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello world",
|
||||||
|
embed_texts=["embed title"],
|
||||||
|
image_urls=["https://example.com/cat.png"],
|
||||||
|
)
|
||||||
|
cache.put(12345, entry)
|
||||||
|
assert cache.pop(12345) is entry
|
||||||
|
|
||||||
|
def test_pop_removes_entry(self) -> None:
|
||||||
|
"""Popping an entry removes it from the cache."""
|
||||||
|
cache = IngestCache()
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
)
|
||||||
|
cache.put(12345, entry)
|
||||||
|
cache.pop(12345)
|
||||||
|
assert cache.pop(12345) is None
|
||||||
|
|
||||||
|
def test_pop_returns_none_for_missing(self) -> None:
|
||||||
|
"""Popping a nonexistent key returns None."""
|
||||||
|
cache = IngestCache()
|
||||||
|
assert cache.pop(99999) is None
|
||||||
|
|
||||||
|
def test_expired_entry_not_returned(self) -> None:
|
||||||
|
"""An expired entry is not returned by pop."""
|
||||||
|
cache = IngestCache(ttl_seconds=0.01)
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
)
|
||||||
|
cache.put(12345, entry)
|
||||||
|
time.sleep(0.02)
|
||||||
|
assert cache.pop(12345) is None
|
||||||
|
|
||||||
|
def test_cleanup_removes_expired(self) -> None:
|
||||||
|
"""Cleanup evicts all expired entries."""
|
||||||
|
cache = IngestCache(ttl_seconds=0.01)
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
)
|
||||||
|
cache.put(1, entry)
|
||||||
|
cache.put(2, entry)
|
||||||
|
time.sleep(0.02)
|
||||||
|
cache._cleanup()
|
||||||
|
assert cache.pop(1) is None
|
||||||
|
assert cache.pop(2) is None
|
||||||
|
|
||||||
|
def test_cleanup_keeps_unexpired(self) -> None:
|
||||||
|
"""Cleanup does not evict entries that are still valid."""
|
||||||
|
cache = IngestCache(ttl_seconds=300)
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
)
|
||||||
|
cache.put(1, entry)
|
||||||
|
cache._cleanup()
|
||||||
|
assert cache.pop(1) is entry
|
||||||
|
|
||||||
|
async def test_start_and_stop(self) -> None:
|
||||||
|
"""The background cleanup task can be started and stopped."""
|
||||||
|
cache = IngestCache(ttl_seconds=0, cleanup_interval_seconds=0.01)
|
||||||
|
entry = CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
)
|
||||||
|
cache.put(1, entry)
|
||||||
|
cache.start()
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
await cache.stop()
|
||||||
|
# Entry should have been cleaned up by the background task.
|
||||||
|
assert cache.pop(1) is None
|
||||||
@@ -0,0 +1,393 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for CLI argument parsing, credentials, and systemd integration.
|
||||||
|
|
||||||
|
Tests cover the CLI helpers and runtime orchestration from crabstero.__main__.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import signal
|
||||||
|
import socket
|
||||||
|
import sys
|
||||||
|
from typing import TYPE_CHECKING, Any, ClassVar, Self
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero import __main__ as cli
|
||||||
|
from crabstero.__main__ import (
|
||||||
|
_parse_args,
|
||||||
|
_read_credential,
|
||||||
|
_sd_notify,
|
||||||
|
_watchdog_interval,
|
||||||
|
_watchdog_loop,
|
||||||
|
)
|
||||||
|
from crabstero.metrics import TcpMetricsAddress, UnixMetricsAddress
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Callable, Coroutine
|
||||||
|
from pathlib import Path
|
||||||
|
from types import TracebackType
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeCrabstero:
|
||||||
|
events: ClassVar[list[object]] = []
|
||||||
|
login_error: ClassVar[BaseException | None] = None
|
||||||
|
close_error: ClassVar[BaseException | None] = None
|
||||||
|
connect_callback: ClassVar[Callable[[], None] | None] = None
|
||||||
|
|
||||||
|
def __init__(self, **kwargs: object) -> None:
|
||||||
|
self._closed = False
|
||||||
|
self.events.append(("init", kwargs))
|
||||||
|
|
||||||
|
async def __aenter__(self) -> Self:
|
||||||
|
self.events.append("enter")
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(
|
||||||
|
self,
|
||||||
|
exc_type: type[BaseException] | None,
|
||||||
|
exc: BaseException | None,
|
||||||
|
traceback: TracebackType | None,
|
||||||
|
) -> None:
|
||||||
|
self.events.append("exit")
|
||||||
|
await self.close()
|
||||||
|
|
||||||
|
async def login(self, token: str) -> None:
|
||||||
|
self.events.append(("login", token))
|
||||||
|
if self.login_error is not None:
|
||||||
|
raise self.login_error
|
||||||
|
|
||||||
|
async def connect(self, **kwargs: object) -> None:
|
||||||
|
self.events.append(("connect", kwargs))
|
||||||
|
callback = type(self).connect_callback
|
||||||
|
if callback is not None:
|
||||||
|
callback()
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
if self._closed:
|
||||||
|
return
|
||||||
|
self.events.append("close")
|
||||||
|
self._closed = True
|
||||||
|
if self.close_error is not None:
|
||||||
|
raise self.close_error
|
||||||
|
|
||||||
|
def is_closed(self) -> bool:
|
||||||
|
return self._closed
|
||||||
|
|
||||||
|
|
||||||
|
def _run_coroutine(coro: Coroutine[Any, Any, None]) -> None:
|
||||||
|
asyncio.run(coro)
|
||||||
|
|
||||||
|
|
||||||
|
def _prepare_main_test(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
*,
|
||||||
|
login_error: BaseException | None = None,
|
||||||
|
close_error: BaseException | None = None,
|
||||||
|
) -> list[object]:
|
||||||
|
_FakeCrabstero.events = []
|
||||||
|
_FakeCrabstero.login_error = login_error
|
||||||
|
_FakeCrabstero.close_error = close_error
|
||||||
|
_FakeCrabstero.connect_callback = None
|
||||||
|
|
||||||
|
def sd_notify(state: str) -> None:
|
||||||
|
_FakeCrabstero.events.append(("notify", state))
|
||||||
|
|
||||||
|
monkeypatch.setattr(sys, "argv", ["crabstero", "--token", "cli-token"])
|
||||||
|
monkeypatch.setattr(cli, "Crabstero", _FakeCrabstero)
|
||||||
|
monkeypatch.setattr("crabstero.__main__.uvloop.run", _run_coroutine)
|
||||||
|
monkeypatch.setattr(cli, "_sd_notify", sd_notify)
|
||||||
|
monkeypatch.setattr(cli, "_watchdog_interval", lambda: None)
|
||||||
|
|
||||||
|
return _FakeCrabstero.events
|
||||||
|
|
||||||
|
|
||||||
|
class TestParseArgs:
|
||||||
|
"""Argument parsing with defaults and overrides."""
|
||||||
|
|
||||||
|
def test_token_from_arg(self) -> None:
|
||||||
|
"""--token flag sets the token."""
|
||||||
|
args = _parse_args(["--token", "abc"])
|
||||||
|
assert args.token == "abc" # noqa: S105
|
||||||
|
|
||||||
|
def test_database_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""Omitting --database-path defaults to 'crabstero.db'."""
|
||||||
|
monkeypatch.delenv("DATABASE_PATH", raising=False)
|
||||||
|
args = _parse_args(["--token", "test"])
|
||||||
|
assert args.database_path == "crabstero.db"
|
||||||
|
|
||||||
|
def test_database_override(self) -> None:
|
||||||
|
"""--database-path overrides the default."""
|
||||||
|
args = _parse_args(["--token", "test", "--database-path", "/custom.db"])
|
||||||
|
assert args.database_path == "/custom.db"
|
||||||
|
|
||||||
|
def test_ingest_only_flag(self) -> None:
|
||||||
|
"""--ingest-only sets ingest_only to True."""
|
||||||
|
args = _parse_args(["--token", "test", "--ingest-only"])
|
||||||
|
assert args.ingest_only is True
|
||||||
|
|
||||||
|
def test_listen_metrics_parses_address(self) -> None:
|
||||||
|
"""--listen-metrics HOST:PORT sets listen_metrics to (host, port)."""
|
||||||
|
args = _parse_args(["--token", "test", "--listen-metrics", "127.0.0.1:9090"])
|
||||||
|
assert args.listen_metrics == TcpMetricsAddress("127.0.0.1", 9090)
|
||||||
|
|
||||||
|
def test_listen_metrics_default_none(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""Omitting --listen-metrics defaults to None."""
|
||||||
|
monkeypatch.delenv("LISTEN_METRICS", raising=False)
|
||||||
|
args = _parse_args(["--token", "test"])
|
||||||
|
assert args.listen_metrics is None
|
||||||
|
|
||||||
|
def test_listen_metrics_ipv6(self) -> None:
|
||||||
|
"""--listen-metrics [::1]:PORT parses IPv6 address correctly."""
|
||||||
|
args = _parse_args(["--token", "test", "--listen-metrics", "[::1]:9090"])
|
||||||
|
assert args.listen_metrics == TcpMetricsAddress("[::1]", 9090)
|
||||||
|
|
||||||
|
def test_listen_metrics_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""LISTEN_METRICS env var is used when --listen-metrics is omitted."""
|
||||||
|
monkeypatch.setenv("LISTEN_METRICS", "127.0.0.1:8080")
|
||||||
|
args = _parse_args(["--token", "test"])
|
||||||
|
assert args.listen_metrics == TcpMetricsAddress("127.0.0.1", 8080)
|
||||||
|
|
||||||
|
def test_listen_metrics_unix_socket(self) -> None:
|
||||||
|
"""--listen-metrics unix:/path sets listen_metrics to a Unix socket."""
|
||||||
|
args = _parse_args(
|
||||||
|
["--token", "test", "--listen-metrics", "unix:/run/crabstero.sock"],
|
||||||
|
)
|
||||||
|
assert args.listen_metrics == UnixMetricsAddress("/run/crabstero.sock")
|
||||||
|
|
||||||
|
def test_listen_metrics_unix_socket_from_env(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""LISTEN_METRICS can enable metrics on a Unix socket."""
|
||||||
|
monkeypatch.setenv("LISTEN_METRICS", "unix:/run/crabstero.sock")
|
||||||
|
args = _parse_args(["--token", "test"])
|
||||||
|
assert args.listen_metrics == UnixMetricsAddress("/run/crabstero.sock")
|
||||||
|
|
||||||
|
def test_listen_metrics_invalid_format(self) -> None:
|
||||||
|
"""--listen-metrics with no colon raises SystemExit."""
|
||||||
|
with pytest.raises(SystemExit):
|
||||||
|
_parse_args(["--token", "test", "--listen-metrics", "bad"])
|
||||||
|
|
||||||
|
def test_listen_metrics_invalid_port(self) -> None:
|
||||||
|
"""--listen-metrics with non-integer port raises SystemExit."""
|
||||||
|
with pytest.raises(SystemExit):
|
||||||
|
_parse_args(["--token", "test", "--listen-metrics", "127.0.0.1:abc"])
|
||||||
|
|
||||||
|
def test_listen_metrics_empty_unix_socket_path(self) -> None:
|
||||||
|
"""--listen-metrics unix: with no path raises SystemExit."""
|
||||||
|
with pytest.raises(SystemExit):
|
||||||
|
_parse_args(["--token", "test", "--listen-metrics", "unix:"])
|
||||||
|
|
||||||
|
def test_missing_token_exits(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""Missing token causes SystemExit."""
|
||||||
|
monkeypatch.delenv("TOKEN", raising=False)
|
||||||
|
monkeypatch.delenv("CREDENTIALS_DIRECTORY", raising=False)
|
||||||
|
with pytest.raises(SystemExit):
|
||||||
|
_parse_args([])
|
||||||
|
|
||||||
|
|
||||||
|
class TestReadCredential:
|
||||||
|
"""Systemd credential file reading."""
|
||||||
|
|
||||||
|
def test_reads_credential_file(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Reads and strips the credential value from the file."""
|
||||||
|
(tmp_path / "mytoken").write_text(" secret123 \n")
|
||||||
|
monkeypatch.setenv("CREDENTIALS_DIRECTORY", str(tmp_path))
|
||||||
|
assert _read_credential("mytoken") == "secret123"
|
||||||
|
|
||||||
|
def test_returns_none_without_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
"""Returns None when CREDENTIALS_DIRECTORY is not set."""
|
||||||
|
monkeypatch.delenv("CREDENTIALS_DIRECTORY", raising=False)
|
||||||
|
assert _read_credential("anything") is None
|
||||||
|
|
||||||
|
def test_returns_none_for_missing_file(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Returns None when the credential file does not exist."""
|
||||||
|
monkeypatch.setenv("CREDENTIALS_DIRECTORY", str(tmp_path))
|
||||||
|
assert _read_credential("nonexistent") is None
|
||||||
|
|
||||||
|
|
||||||
|
class TestSystemdNotify:
|
||||||
|
"""Systemd notification helper behavior."""
|
||||||
|
|
||||||
|
def test_sd_notify_noop_without_socket(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Missing NOTIFY_SOCKET is a no-op."""
|
||||||
|
monkeypatch.delenv("NOTIFY_SOCKET", raising=False)
|
||||||
|
_sd_notify("READY=1")
|
||||||
|
|
||||||
|
def test_sd_notify_sends_datagram(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""_sd_notify sends the payload to NOTIFY_SOCKET."""
|
||||||
|
socket_path = tmp_path / "notify.sock"
|
||||||
|
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as server:
|
||||||
|
server.bind(str(socket_path))
|
||||||
|
server.settimeout(1)
|
||||||
|
monkeypatch.setenv("NOTIFY_SOCKET", str(socket_path))
|
||||||
|
|
||||||
|
_sd_notify("READY=1")
|
||||||
|
|
||||||
|
assert server.recv(1024) == b"READY=1"
|
||||||
|
|
||||||
|
|
||||||
|
class TestWatchdog:
|
||||||
|
"""Systemd watchdog helper behavior."""
|
||||||
|
|
||||||
|
def test_watchdog_interval_missing(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Missing WATCHDOG_USEC disables watchdog pings."""
|
||||||
|
monkeypatch.delenv("WATCHDOG_USEC", raising=False)
|
||||||
|
assert _watchdog_interval() is None
|
||||||
|
|
||||||
|
def test_watchdog_interval_valid(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""WATCHDOG_USEC is converted to half the interval in seconds."""
|
||||||
|
monkeypatch.setenv("WATCHDOG_USEC", "4000000")
|
||||||
|
assert _watchdog_interval() == 2
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"value",
|
||||||
|
[
|
||||||
|
pytest.param("invalid", id="not-integer"),
|
||||||
|
pytest.param("0", id="zero"),
|
||||||
|
pytest.param("-1", id="negative"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_watchdog_interval_invalid(
|
||||||
|
self,
|
||||||
|
value: str,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Invalid WATCHDOG_USEC values disable watchdog pings."""
|
||||||
|
monkeypatch.setenv("WATCHDOG_USEC", value)
|
||||||
|
assert _watchdog_interval() is None
|
||||||
|
|
||||||
|
async def test_watchdog_loop_sends_until_stopped(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""_watchdog_loop sends WATCHDOG=1 and exits when stopped."""
|
||||||
|
notifications: list[str] = []
|
||||||
|
stop = asyncio.Event()
|
||||||
|
monkeypatch.setattr(cli, "_sd_notify", notifications.append)
|
||||||
|
|
||||||
|
task = asyncio.create_task(_watchdog_loop(60, stop))
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
stop.set()
|
||||||
|
await asyncio.wait_for(task, timeout=1)
|
||||||
|
|
||||||
|
assert notifications == ["WATCHDOG=1"]
|
||||||
|
|
||||||
|
|
||||||
|
class TestMainLifecycle:
|
||||||
|
"""CLI runtime orchestration."""
|
||||||
|
|
||||||
|
def test_ready_sent_after_login_before_connect(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""READY=1 is sent after login succeeds and before connect starts."""
|
||||||
|
events = _prepare_main_test(monkeypatch)
|
||||||
|
|
||||||
|
cli.main()
|
||||||
|
|
||||||
|
assert events.index(("login", "cli-token")) < events.index(
|
||||||
|
("notify", "READY=1"),
|
||||||
|
)
|
||||||
|
assert events.index(("notify", "READY=1")) < events.index(
|
||||||
|
("connect", {"reconnect": True}),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_stopping_sent_after_clean_connect_return(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""STOPPING=1 is sent when the CLI run exits cleanly after readiness."""
|
||||||
|
events = _prepare_main_test(monkeypatch)
|
||||||
|
|
||||||
|
cli.main()
|
||||||
|
|
||||||
|
assert events.index(("connect", {"reconnect": True})) < events.index(
|
||||||
|
("notify", "STOPPING=1"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_ready_not_sent_when_login_fails(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""A login failure exits without reporting readiness."""
|
||||||
|
events = _prepare_main_test(
|
||||||
|
monkeypatch,
|
||||||
|
login_error=RuntimeError("login failed"),
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="login failed"):
|
||||||
|
cli.main()
|
||||||
|
|
||||||
|
assert ("notify", "READY=1") not in events
|
||||||
|
assert ("notify", "STOPPING=1") not in events
|
||||||
|
|
||||||
|
def test_shutdown_task_exception_is_propagated(
|
||||||
|
self,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""A completed shutdown task exception is not silently dropped."""
|
||||||
|
signal_handlers: dict[int, Callable[[], None]] = {}
|
||||||
|
|
||||||
|
class SignalLoop:
|
||||||
|
def add_signal_handler(
|
||||||
|
self,
|
||||||
|
sig: int,
|
||||||
|
callback: Callable[[], None],
|
||||||
|
) -> None:
|
||||||
|
signal_handlers[sig] = callback
|
||||||
|
|
||||||
|
def remove_signal_handler(self, sig: int) -> None:
|
||||||
|
signal_handlers.pop(sig, None)
|
||||||
|
|
||||||
|
events = _prepare_main_test(
|
||||||
|
monkeypatch,
|
||||||
|
close_error=RuntimeError("close failed"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def request_shutdown() -> None:
|
||||||
|
signal_handlers[signal.SIGTERM]()
|
||||||
|
|
||||||
|
monkeypatch.setattr(asyncio, "get_running_loop", SignalLoop)
|
||||||
|
_FakeCrabstero.connect_callback = request_shutdown
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="close failed"):
|
||||||
|
cli.main()
|
||||||
|
|
||||||
|
assert events.index(("notify", "STOPPING=1")) < events.index("close")
|
||||||
@@ -0,0 +1,641 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for the Database class.
|
||||||
|
|
||||||
|
Tests cover Database.connect (pragmas, schema), markov start word and
|
||||||
|
transition CRUD, image storage, flag CRUD, and channel ingestion tracking.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero.database import ChannelImage, Database, StartWord, Transition
|
||||||
|
from crabstero.flags import EntityType, Flag
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class TestConnect:
|
||||||
|
"""Database.connect creates a configured SQLite database."""
|
||||||
|
|
||||||
|
async def test_synchronous_normal(self, db: Database) -> None:
|
||||||
|
"""Synchronous mode is set to NORMAL."""
|
||||||
|
async with db._connection.execute("PRAGMA synchronous") as cursor:
|
||||||
|
row = await cursor.fetchone()
|
||||||
|
assert row is not None
|
||||||
|
assert row[0] == 1
|
||||||
|
|
||||||
|
async def test_schema_creates_tables(self, db: Database) -> None:
|
||||||
|
"""All expected tables exist after connect."""
|
||||||
|
async with db._connection.execute(
|
||||||
|
"SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name",
|
||||||
|
) as cursor:
|
||||||
|
tables = [row[0] for row in await cursor.fetchall()]
|
||||||
|
assert tables == [
|
||||||
|
"channel_images",
|
||||||
|
"flags",
|
||||||
|
"ingested_channels",
|
||||||
|
"markov_start_words",
|
||||||
|
"markov_transitions",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class TestAddMarkovData:
|
||||||
|
"""Markov start word and transition storage via add_markov_data."""
|
||||||
|
|
||||||
|
async def test_stores_start_word(self, db: Database) -> None:
|
||||||
|
"""Inserted start word can be retrieved by channel."""
|
||||||
|
await db.add_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
result = await db.get_random_start_word(1)
|
||||||
|
assert result == "Hello"
|
||||||
|
|
||||||
|
async def test_stores_transition(self, db: Database) -> None:
|
||||||
|
"""Inserted transition can be retrieved by channel and word."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == "world."
|
||||||
|
|
||||||
|
async def test_stores_both_in_single_call(self, db: Database) -> None:
|
||||||
|
"""Start words and transitions are stored in a single call."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello")],
|
||||||
|
[Transition(1, 100, "Hello", "world.")],
|
||||||
|
)
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_empty_lists_is_noop(self, db: Database) -> None:
|
||||||
|
"""Empty lists do not store any data."""
|
||||||
|
await db.add_markov_data([], [])
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
|
||||||
|
class TestMarkovReadMethods:
|
||||||
|
"""Markov read methods return None for out-of-scope or missing data."""
|
||||||
|
|
||||||
|
async def test_returns_none_when_empty(self, db: Database) -> None:
|
||||||
|
"""Returns None when no data has been stored."""
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "nonexistent") is None
|
||||||
|
assert await db.get_random_completing_next_word(1, "nonexistent") is None
|
||||||
|
|
||||||
|
async def test_start_word_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""A start word in one channel is not returned for another channel."""
|
||||||
|
await db.add_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
assert await db.get_random_start_word(2) is None
|
||||||
|
|
||||||
|
async def test_next_word_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""A transition in one channel is not returned for another channel."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(2, "Hello") is None
|
||||||
|
|
||||||
|
async def test_completing_next_word_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""Completing transition is not returned for another channel."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_completing_next_word(2, "Hello") is None
|
||||||
|
|
||||||
|
async def test_next_word_scoped_to_word(self, db: Database) -> None:
|
||||||
|
"""A transition for one word is not returned when querying a different word."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Goodbye") is None
|
||||||
|
|
||||||
|
async def test_completing_next_word_scoped_to_word(self, db: Database) -> None:
|
||||||
|
"""Completing transition is not returned for a different word."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_completing_next_word(1, "Goodbye") is None
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"completing_word",
|
||||||
|
[
|
||||||
|
pytest.param("world.", id="period"),
|
||||||
|
pytest.param("world!", id="exclamation"),
|
||||||
|
pytest.param("world?", id="question"),
|
||||||
|
pytest.param("world\u00a7", id="section-sign"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_completing_word_filters_punctuation(
|
||||||
|
self,
|
||||||
|
db: Database,
|
||||||
|
completing_word: str,
|
||||||
|
) -> None:
|
||||||
|
"""get_random_completing_next_word only returns sentence-ending words."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "beautiful"),
|
||||||
|
Transition(1, 100, "Hello", completing_word),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
for _ in range(100):
|
||||||
|
result = await db.get_random_completing_next_word(1, "Hello")
|
||||||
|
assert result == completing_word
|
||||||
|
|
||||||
|
async def test_completing_returns_none_without_match(self, db: Database) -> None:
|
||||||
|
"""Returns None when no transitions end with sentence punctuation."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "beautiful")])
|
||||||
|
result = await db.get_random_completing_next_word(1, "Hello")
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
async def test_start_word_pooled_across_users(self, db: Database) -> None:
|
||||||
|
"""Start words from different users are visible in the same channel query."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello"), StartWord(1, 200, "Goodbye")],
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
assert await db.get_random_start_word(1) in {"Hello", "Goodbye"}
|
||||||
|
|
||||||
|
async def test_next_word_pooled_across_users(self, db: Database) -> None:
|
||||||
|
"""Transitions from different users are visible in the same channel query."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 200, "Hello", "friend."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
assert await db.get_random_next_word(1, "Hello") in {"world.", "friend."}
|
||||||
|
|
||||||
|
async def test_completing_next_word_pooled_across_users(self, db: Database) -> None:
|
||||||
|
"""Completing transitions from different users are visible."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 200, "Hello", "friend."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
assert await db.get_random_completing_next_word(1, "Hello") in {
|
||||||
|
"world.",
|
||||||
|
"friend.",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestRemoveMarkovData:
|
||||||
|
"""Markov data removal via remove_markov_data."""
|
||||||
|
|
||||||
|
async def test_removes_one_start_word(self, db: Database) -> None:
|
||||||
|
"""Removes exactly one matching start word row."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello"), StartWord(1, 100, "Hello")],
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
# One copy should remain.
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
|
||||||
|
async def test_removes_one_transition(self, db: Database) -> None:
|
||||||
|
"""Removes exactly one matching transition row."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_removes_last_start_word(self, db: Database) -> None:
|
||||||
|
"""Removing the only start word leaves the table empty for that channel."""
|
||||||
|
await db.add_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_removes_last_transition(self, db: Database) -> None:
|
||||||
|
"""Removing the only transition leaves no next word."""
|
||||||
|
await db.add_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
|
||||||
|
async def test_no_match_is_noop(self, db: Database) -> None:
|
||||||
|
"""Removing a non-existent row does not raise."""
|
||||||
|
await db.remove_markov_data(
|
||||||
|
[StartWord(1, 100, "nope")],
|
||||||
|
[Transition(1, 100, "nope", "nah")],
|
||||||
|
)
|
||||||
|
|
||||||
|
async def test_start_word_removal_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""Removing a start word in one channel leaves another channel intact."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello"), StartWord(2, 200, "Hello")],
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_start_word(2) == "Hello"
|
||||||
|
|
||||||
|
async def test_transition_removal_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""Removing a transition in one channel leaves another channel intact."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(2, 200, "Hello", "world."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
assert await db.get_random_next_word(2, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_start_word_removal_scoped_to_user(self, db: Database) -> None:
|
||||||
|
"""Removing a start word for one user leaves another user."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello"), StartWord(1, 200, "Hello")],
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
|
||||||
|
async def test_transition_removal_scoped_to_user(self, db: Database) -> None:
|
||||||
|
"""Removing a transition for one user leaves another user."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 200, "Hello", "world."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_start_word_removal_scoped_to_word(self, db: Database) -> None:
|
||||||
|
"""Removing one start word leaves a different start word."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello"), StartWord(1, 100, "World")],
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
assert await db.get_random_start_word(1) == "World"
|
||||||
|
|
||||||
|
async def test_transition_removal_scoped_to_word(self, db: Database) -> None:
|
||||||
|
"""Removing one word's transition leaves another word's."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 100, "Goodbye", "world."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
assert await db.get_random_next_word(1, "Goodbye") == "world."
|
||||||
|
|
||||||
|
async def test_transition_removal_scoped_to_next_word(self, db: Database) -> None:
|
||||||
|
"""Removing one next_word leaves a different next_word."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Hello", "world."),
|
||||||
|
Transition(1, 100, "Hello", "friend."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_markov_data([], [Transition(1, 100, "Hello", "world.")])
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "friend."
|
||||||
|
|
||||||
|
async def test_empty_lists_is_noop(self, db: Database) -> None:
|
||||||
|
"""Empty lists do not error."""
|
||||||
|
await db.remove_markov_data([], [])
|
||||||
|
|
||||||
|
|
||||||
|
class TestImages:
|
||||||
|
"""Image URL storage and random retrieval."""
|
||||||
|
|
||||||
|
async def test_add_and_retrieve(self, db: Database) -> None:
|
||||||
|
"""Inserted image URL can be retrieved by channel."""
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/cat.png")])
|
||||||
|
result = await db.get_random_image(1)
|
||||||
|
assert result == "https://example.com/cat.png"
|
||||||
|
|
||||||
|
async def test_returns_none_when_empty(self, db: Database) -> None:
|
||||||
|
"""Returns None for a channel with no images."""
|
||||||
|
result = await db.get_random_image(999)
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
async def test_image_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""An image in one channel is not returned for another channel."""
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/cat.png")])
|
||||||
|
assert await db.get_random_image(2) is None
|
||||||
|
|
||||||
|
async def test_empty_list_is_noop(self, db: Database) -> None:
|
||||||
|
"""Empty list does not store any data."""
|
||||||
|
await db.add_images([])
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
|
||||||
|
async def test_image_pooled_across_users(self, db: Database) -> None:
|
||||||
|
"""Images from different users are visible in the same channel query."""
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(1, 200, "https://example.com/b.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
assert await db.get_random_image(1) in {
|
||||||
|
"https://example.com/a.png",
|
||||||
|
"https://example.com/b.png",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestRemoveImages:
|
||||||
|
"""Image removal via remove_images."""
|
||||||
|
|
||||||
|
async def test_removes_one_image(self, db: Database) -> None:
|
||||||
|
"""Removes exactly one matching image row."""
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
# One copy should remain.
|
||||||
|
assert await db.get_random_image(1) == "https://example.com/a.png"
|
||||||
|
|
||||||
|
async def test_removes_last_image(self, db: Database) -> None:
|
||||||
|
"""Removing the only image leaves none for that channel."""
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
|
||||||
|
async def test_image_removal_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""Removing an image in one channel leaves another channel intact."""
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(2, 200, "https://example.com/a.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
assert await db.get_random_image(2) == "https://example.com/a.png"
|
||||||
|
|
||||||
|
async def test_image_removal_scoped_to_user(self, db: Database) -> None:
|
||||||
|
"""Removing an image for one user leaves another user."""
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(1, 200, "https://example.com/a.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
assert await db.get_random_image(1) == "https://example.com/a.png"
|
||||||
|
|
||||||
|
async def test_image_removal_scoped_to_url(self, db: Database) -> None:
|
||||||
|
"""Removing one URL leaves a different URL for the same user."""
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(1, 100, "https://example.com/b.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
assert await db.get_random_image(1) == "https://example.com/b.png"
|
||||||
|
|
||||||
|
async def test_no_match_is_noop(self, db: Database) -> None:
|
||||||
|
"""Removing a non-existent image does not raise."""
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/nope.png")])
|
||||||
|
|
||||||
|
async def test_empty_list_is_noop(self, db: Database) -> None:
|
||||||
|
"""Empty list does not error."""
|
||||||
|
await db.remove_images([])
|
||||||
|
|
||||||
|
|
||||||
|
class TestFlags:
|
||||||
|
"""Flag CRUD operations on entities."""
|
||||||
|
|
||||||
|
async def test_set_and_check(self, db: Database) -> None:
|
||||||
|
"""A set flag is reported as set."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_REPLY) is True
|
||||||
|
|
||||||
|
async def test_unset_flag_is_false(self, db: Database) -> None:
|
||||||
|
"""An unset flag is reported as not set."""
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_clear_flag(self, db: Database) -> None:
|
||||||
|
"""A cleared flag is no longer reported as set."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.clear_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_set_idempotent(self, db: Database) -> None:
|
||||||
|
"""Setting the same flag twice does not raise."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_REPLY) is True
|
||||||
|
|
||||||
|
async def test_scoped_to_entity_id(self, db: Database) -> None:
|
||||||
|
"""A flag set on one entity is not visible on another entity."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "456", Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_scoped_to_entity_type(self, db: Database) -> None:
|
||||||
|
"""A flag set on one entity type is not visible on another."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "123", Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_scoped_to_flag_name(self, db: Database) -> None:
|
||||||
|
"""A flag set under one name is not visible under a different name."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_INGEST) is False
|
||||||
|
|
||||||
|
async def test_clear_scoped_to_flag_name(self, db: Database) -> None:
|
||||||
|
"""Clearing one flag leaves other flags on the same entity intact."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_INGEST)
|
||||||
|
await db.clear_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_INGEST) is True
|
||||||
|
|
||||||
|
async def test_clear_scoped_to_entity_id(self, db: Database) -> None:
|
||||||
|
"""Clearing a flag on one entity leaves another entity."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "456", Flag.NO_REPLY)
|
||||||
|
await db.clear_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "456", Flag.NO_REPLY) is True
|
||||||
|
|
||||||
|
async def test_clear_scoped_to_entity_type(self, db: Database) -> None:
|
||||||
|
"""Clearing a flag on one entity type leaves another type."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.USER, "123", Flag.NO_REPLY)
|
||||||
|
await db.clear_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "123", Flag.NO_REPLY) is True
|
||||||
|
|
||||||
|
async def test_clear_unset_flag_is_noop(self, db: Database) -> None:
|
||||||
|
"""Clearing a flag that was never set does not raise or affect other flags."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "123", Flag.NO_REPLY)
|
||||||
|
await db.clear_flag(EntityType.CHANNEL, "123", Flag.NO_INGEST)
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "123", Flag.NO_REPLY) is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestChannelIngestion:
|
||||||
|
"""Channel ingestion tracking."""
|
||||||
|
|
||||||
|
async def test_mark_and_check(self, db: Database) -> None:
|
||||||
|
"""A marked channel is reported as ingested."""
|
||||||
|
await db.mark_channel_ingested(42)
|
||||||
|
assert await db.is_channel_ingested(42) is True
|
||||||
|
|
||||||
|
async def test_not_ingested_by_default(self, db: Database) -> None:
|
||||||
|
"""Unmarked channels are not reported as ingested."""
|
||||||
|
assert await db.is_channel_ingested(42) is False
|
||||||
|
|
||||||
|
async def test_mark_idempotent(self, db: Database) -> None:
|
||||||
|
"""Marking the same channel twice does not raise."""
|
||||||
|
await db.mark_channel_ingested(42)
|
||||||
|
await db.mark_channel_ingested(42)
|
||||||
|
assert await db.is_channel_ingested(42) is True
|
||||||
|
|
||||||
|
async def test_scoped_to_channel(self, db: Database) -> None:
|
||||||
|
"""Marking one channel as ingested does not affect another channel."""
|
||||||
|
await db.mark_channel_ingested(42)
|
||||||
|
assert await db.is_channel_ingested(99) is False
|
||||||
|
|
||||||
|
|
||||||
|
class TestTransactionRollback:
|
||||||
|
"""Transaction rolls back all changes on error."""
|
||||||
|
|
||||||
|
async def test_error_rolls_back_insert(self, db: Database) -> None:
|
||||||
|
"""An error during a transaction prevents partial data from persisting."""
|
||||||
|
with pytest.raises(RuntimeError, match="simulated"):
|
||||||
|
await self._insert_and_fail(db)
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _insert_and_fail(db: Database) -> None:
|
||||||
|
async with db._transaction():
|
||||||
|
await db._connection.execute(
|
||||||
|
"INSERT INTO markov_start_words"
|
||||||
|
" (channel_id, user_id, word)"
|
||||||
|
" VALUES (?, ?, ?)",
|
||||||
|
(1, 100, "should_not_persist"),
|
||||||
|
)
|
||||||
|
msg = "simulated failure"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
|
||||||
|
|
||||||
|
class TestForgetUser:
|
||||||
|
"""Atomic forget-user transaction across all tables."""
|
||||||
|
|
||||||
|
async def test_deletes_data_and_sets_no_ingest(self, db: Database) -> None:
|
||||||
|
"""All user data is removed and noIngest flag is set."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello")],
|
||||||
|
[Transition(1, 100, "Hello", "world.")],
|
||||||
|
)
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/a.png")])
|
||||||
|
await db.set_flag(EntityType.USER, "100", Flag.ALLOW_PINGS)
|
||||||
|
|
||||||
|
await db.forget_user(100, Flag.NO_INGEST)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.ALLOW_PINGS) is False
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.NO_INGEST) is True
|
||||||
|
|
||||||
|
async def test_preserves_other_users(self, db: Database) -> None:
|
||||||
|
"""Data belonging to other users is not affected."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Gone"), StartWord(1, 200, "Keep")],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "Gone", "away."),
|
||||||
|
Transition(1, 200, "Keep", "this."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/gone.png"),
|
||||||
|
ChannelImage(1, 200, "https://example.com/stay.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.set_flag(EntityType.USER, "200", Flag.ALLOW_PINGS)
|
||||||
|
|
||||||
|
await db.forget_user(100, Flag.NO_INGEST)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) == "Keep"
|
||||||
|
assert await db.get_random_next_word(1, "Keep") == "this."
|
||||||
|
assert await db.get_random_image(1) == "https://example.com/stay.png"
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "200", Flag.ALLOW_PINGS) is True
|
||||||
|
|
||||||
|
async def test_preserves_other_entity_type_flags(self, db: Database) -> None:
|
||||||
|
"""Flags on channels with the same entity ID are not affected."""
|
||||||
|
await db.set_flag(EntityType.CHANNEL, "100", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.USER, "100", Flag.NO_REPLY)
|
||||||
|
|
||||||
|
await db.forget_user(100, Flag.NO_INGEST)
|
||||||
|
|
||||||
|
assert await db.is_flag_set(EntityType.CHANNEL, "100", Flag.NO_REPLY) is True
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_clears_existing_flags_except_no_ingest(self, db: Database) -> None:
|
||||||
|
"""Existing user flags are cleared but noIngest remains."""
|
||||||
|
await db.set_flag(EntityType.USER, "100", Flag.NO_REPLY)
|
||||||
|
await db.set_flag(EntityType.USER, "100", Flag.ALLOW_PINGS)
|
||||||
|
|
||||||
|
await db.forget_user(100, Flag.NO_INGEST)
|
||||||
|
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.NO_REPLY) is False
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.ALLOW_PINGS) is False
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "100", Flag.NO_INGEST) is True
|
||||||
|
|
||||||
|
async def test_deletes_across_channels(self, db: Database) -> None:
|
||||||
|
"""All user data is removed from every channel."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "One"), StartWord(2, 100, "Two")],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "One", "fish."),
|
||||||
|
Transition(2, 100, "Two", "fish."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await db.add_images(
|
||||||
|
[
|
||||||
|
ChannelImage(1, 100, "https://example.com/a.png"),
|
||||||
|
ChannelImage(2, 100, "https://example.com/b.png"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
await db.forget_user(100, Flag.NO_INGEST)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_start_word(2) is None
|
||||||
|
assert await db.get_random_next_word(1, "One") is None
|
||||||
|
assert await db.get_random_next_word(2, "Two") is None
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
assert await db.get_random_image(2) is None
|
||||||
|
|
||||||
|
async def test_noop_for_nonexistent_user(self, db: Database) -> None:
|
||||||
|
"""Forgetting a user with no data does not raise."""
|
||||||
|
await db.forget_user(999, Flag.NO_INGEST)
|
||||||
|
assert await db.is_flag_set(EntityType.USER, "999", Flag.NO_INGEST) is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestWriteDurability:
|
||||||
|
"""Writes persist across close and reopen."""
|
||||||
|
|
||||||
|
async def test_markov_data_survives_reopen(self, tmp_path: Path) -> None:
|
||||||
|
"""Data written via add_markov_data is durable after close/reopen."""
|
||||||
|
db_path = str(tmp_path / "durability.db")
|
||||||
|
db = await Database.connect(db_path)
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "Hello")],
|
||||||
|
[Transition(1, 100, "Hello", "world.")],
|
||||||
|
)
|
||||||
|
await db.close()
|
||||||
|
|
||||||
|
db2 = await Database.connect(db_path)
|
||||||
|
assert await db2.get_random_start_word(1) == "Hello"
|
||||||
|
assert await db2.get_random_next_word(1, "Hello") == "world."
|
||||||
|
await db2.close()
|
||||||
@@ -0,0 +1,116 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for the flag convenience wrappers and enums.
|
||||||
|
|
||||||
|
Tests cover _entity_id, Flag and EntityType enums, and the high-level
|
||||||
|
set_flag/clear_flag/is_flag_set wrappers from crabstero.flags.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero.flags import (
|
||||||
|
EntityType,
|
||||||
|
Flag,
|
||||||
|
_entity_id,
|
||||||
|
clear_flag,
|
||||||
|
is_flag_set,
|
||||||
|
set_flag,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from crabstero.database import Database
|
||||||
|
|
||||||
|
|
||||||
|
class _StubSnowflake:
|
||||||
|
"""Minimal stand-in for discord.abc.Snowflake."""
|
||||||
|
|
||||||
|
def __init__(self, *, entity_id: int = 123456789) -> None:
|
||||||
|
self.id = entity_id
|
||||||
|
|
||||||
|
|
||||||
|
class TestEntityId:
|
||||||
|
"""Entity ID extraction from Discord objects and integers."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("entity", "expected"),
|
||||||
|
[
|
||||||
|
pytest.param(42, "42", id="raw-integer"),
|
||||||
|
pytest.param(
|
||||||
|
_StubSnowflake(entity_id=99),
|
||||||
|
"99",
|
||||||
|
id="snowflake-object",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_extraction(self, entity: object, expected: str) -> None:
|
||||||
|
"""Converts the entity to the expected string ID."""
|
||||||
|
assert _entity_id(entity) == expected # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
|
class TestFlagEnums:
|
||||||
|
"""Flag and EntityType enum values."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("member", "expected"),
|
||||||
|
[
|
||||||
|
pytest.param(Flag.NO_REPLY, "noReply", id="no-reply"),
|
||||||
|
pytest.param(Flag.NO_INGEST, "noIngest", id="no-ingest"),
|
||||||
|
pytest.param(Flag.ALLOW_PINGS, "allowPings", id="allow-pings"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_flag_values(self, member: Flag, expected: str) -> None:
|
||||||
|
"""Flag enum value matches the expected database string."""
|
||||||
|
assert member.value == expected
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("member", "expected"),
|
||||||
|
[
|
||||||
|
pytest.param(EntityType.CHANNEL, "channel", id="channel"),
|
||||||
|
pytest.param(EntityType.SERVER, "server", id="server"),
|
||||||
|
pytest.param(EntityType.USER, "user", id="user"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_entity_type_values(self, member: EntityType, expected: str) -> None:
|
||||||
|
"""EntityType enum value matches the expected database string."""
|
||||||
|
assert member.value == expected
|
||||||
|
|
||||||
|
|
||||||
|
class TestSetClearCheck:
|
||||||
|
"""High-level flag set/clear/check cycle through the flags module."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"flag",
|
||||||
|
[
|
||||||
|
pytest.param(Flag.NO_REPLY, id="no-reply"),
|
||||||
|
pytest.param(Flag.NO_INGEST, id="no-ingest"),
|
||||||
|
pytest.param(Flag.ALLOW_PINGS, id="allow-pings"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_set_then_check(self, db: Database, flag: Flag) -> None:
|
||||||
|
"""A set flag is reported as set."""
|
||||||
|
await set_flag(db, 1, EntityType.CHANNEL, flag)
|
||||||
|
assert await is_flag_set(db, 1, EntityType.CHANNEL, flag) is True
|
||||||
|
|
||||||
|
async def test_unset_returns_false(self, db: Database) -> None:
|
||||||
|
"""An unset flag is reported as not set."""
|
||||||
|
assert await is_flag_set(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) is False
|
||||||
|
|
||||||
|
async def test_clear_removes_flag(self, db: Database) -> None:
|
||||||
|
"""A cleared flag is no longer reported as set."""
|
||||||
|
await set_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY)
|
||||||
|
await clear_flag(db, 1, EntityType.CHANNEL, Flag.NO_REPLY)
|
||||||
|
assert await is_flag_set(db, 1, EntityType.CHANNEL, Flag.NO_REPLY) is False
|
||||||
@@ -0,0 +1,348 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for Markov chain ingestion and generation.
|
||||||
|
|
||||||
|
Tests cover is_complete_sentence, _ingest_sentence, ingest, and generate
|
||||||
|
from crabstero.markov.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero.database import StartWord, Transition
|
||||||
|
from crabstero.markov import (
|
||||||
|
DEFAULT_SENTENCE_END,
|
||||||
|
_ingest_sentence,
|
||||||
|
_uningest_sentence,
|
||||||
|
generate,
|
||||||
|
ingest,
|
||||||
|
is_complete_sentence,
|
||||||
|
uningest,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from crabstero.database import Database
|
||||||
|
|
||||||
|
|
||||||
|
class TestIsCompleteSentence:
|
||||||
|
"""Sentence completeness detection."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("sentence", "expected"),
|
||||||
|
[
|
||||||
|
pytest.param("Hello.", True, id="period"),
|
||||||
|
pytest.param("Wow!", True, id="exclamation"),
|
||||||
|
pytest.param("Really?", True, id="question"),
|
||||||
|
pytest.param(f"Hello{DEFAULT_SENTENCE_END}", True, id="section-sign"),
|
||||||
|
pytest.param("Hello", False, id="no-punctuation"),
|
||||||
|
pytest.param("", False, id="empty-string"),
|
||||||
|
pytest.param("Hello. ", False, id="trailing-space"),
|
||||||
|
pytest.param("Hello,", False, id="comma"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_detection(self, sentence: str, expected: bool) -> None: # noqa: FBT001
|
||||||
|
"""Correctly identifies sentence completeness."""
|
||||||
|
assert is_complete_sentence(sentence) is expected
|
||||||
|
|
||||||
|
|
||||||
|
class TestIngestSentence:
|
||||||
|
"""Single sentence ingestion into the Markov chain."""
|
||||||
|
|
||||||
|
async def test_stores_start_word(self, db: Database) -> None:
|
||||||
|
"""First word of the sentence is stored as a start word."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
result = await db.get_random_start_word(1)
|
||||||
|
assert result == "Hello"
|
||||||
|
|
||||||
|
async def test_stores_transitions(self, db: Database) -> None:
|
||||||
|
"""Adjacent words create transitions."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == "world."
|
||||||
|
|
||||||
|
async def test_appends_sentence_end_if_missing(self, db: Database) -> None:
|
||||||
|
"""Unpunctuated sentence gets the default sentence-end marker."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world")
|
||||||
|
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == f"world{DEFAULT_SENTENCE_END}"
|
||||||
|
|
||||||
|
async def test_preserves_existing_punctuation(self, db: Database) -> None:
|
||||||
|
"""Already-punctuated sentence keeps its terminator."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world!")
|
||||||
|
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == "world!"
|
||||||
|
|
||||||
|
async def test_single_word_stores_nothing(self, db: Database) -> None:
|
||||||
|
"""A single-word sentence produces no start words or transitions."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello.")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "Hello.") is None
|
||||||
|
|
||||||
|
async def test_stores_all_transitions(self, db: Database) -> None:
|
||||||
|
"""All adjacent word pairs create transitions."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "A B C.")
|
||||||
|
|
||||||
|
assert await db.get_random_next_word(1, "A") == "B"
|
||||||
|
assert await db.get_random_next_word(1, "B") == "C."
|
||||||
|
|
||||||
|
|
||||||
|
class TestIngestParagraph:
|
||||||
|
"""Paragraph ingestion splits into sentences."""
|
||||||
|
|
||||||
|
async def test_single_sentence(self, db: Database) -> None:
|
||||||
|
"""A single sentence paragraph is ingested."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
|
||||||
|
async def test_multiple_sentences(self, db: Database) -> None:
|
||||||
|
"""Multiple sentences are split and ingested individually."""
|
||||||
|
await ingest(db, 1, 100, "Hello world. Goodbye world!")
|
||||||
|
|
||||||
|
# Both "Hello" and "Goodbye" should appear as start words.
|
||||||
|
async with db._connection.execute(
|
||||||
|
"SELECT DISTINCT word FROM markov_start_words WHERE channel_id = ?",
|
||||||
|
(1,),
|
||||||
|
) as cursor:
|
||||||
|
start_words = {row[0] for row in await cursor.fetchall()}
|
||||||
|
assert start_words == {"Hello", "Goodbye"}
|
||||||
|
|
||||||
|
async def test_normalizes_whitespace(self, db: Database) -> None:
|
||||||
|
"""Extra spaces and newlines are collapsed."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.\nGoodbye world!")
|
||||||
|
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == "world."
|
||||||
|
|
||||||
|
async def test_appends_default_end(self, db: Database) -> None:
|
||||||
|
"""Unpunctuated paragraph gets the default sentence-end marker."""
|
||||||
|
await ingest(db, 1, 100, "Hello world")
|
||||||
|
|
||||||
|
result = await db.get_random_next_word(1, "Hello")
|
||||||
|
assert result == f"world{DEFAULT_SENTENCE_END}"
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"paragraph",
|
||||||
|
[
|
||||||
|
pytest.param("", id="empty-string"),
|
||||||
|
pytest.param(" ", id="whitespace-only"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_empty_input_stores_nothing(
|
||||||
|
self,
|
||||||
|
db: Database,
|
||||||
|
paragraph: str,
|
||||||
|
) -> None:
|
||||||
|
"""Empty or whitespace-only input does not store any data."""
|
||||||
|
await ingest(db, 1, 100, paragraph)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_splits_on_punctuation_followed_by_space(self, db: Database) -> None:
|
||||||
|
"""Punctuation followed by a space splits into separate sentences."""
|
||||||
|
await ingest(db, 1, 100, "Dr. Smith likes cats")
|
||||||
|
|
||||||
|
# "Dr." splits off as a single-word sentence (stores nothing).
|
||||||
|
# "Smith likes cats" becomes a sentence with "Smith" as start word.
|
||||||
|
assert await db.get_random_start_word(1) == "Smith"
|
||||||
|
|
||||||
|
async def test_no_split_without_space_after_punctuation(self, db: Database) -> None:
|
||||||
|
"""Punctuation not followed by a space keeps words together."""
|
||||||
|
await ingest(db, 1, 100, "Hello.World is here")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) == "Hello.World"
|
||||||
|
|
||||||
|
|
||||||
|
class TestGenerate:
|
||||||
|
"""Markov chain text generation."""
|
||||||
|
|
||||||
|
async def test_fallback_on_empty_channel(self, db: Database) -> None:
|
||||||
|
"""Returns an informational message when the channel has no data."""
|
||||||
|
result = await generate(db, 1)
|
||||||
|
expected = (
|
||||||
|
"I do not have enough data to generate a message yet."
|
||||||
|
" Chat a bit more so I can learn how this channel talks."
|
||||||
|
)
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
async def test_generates_from_ingested_data(self, db: Database) -> None:
|
||||||
|
"""Generated text uses words from ingested data."""
|
||||||
|
await ingest(db, 1, 100, "The quick brown fox.")
|
||||||
|
|
||||||
|
result = await generate(db, 1)
|
||||||
|
assert result == "The quick brown fox."
|
||||||
|
|
||||||
|
async def test_strips_section_sign(self, db: Database) -> None:
|
||||||
|
"""The internal section sign marker never appears in output."""
|
||||||
|
await ingest(db, 1, 100, "Hello world")
|
||||||
|
|
||||||
|
result = await generate(db, 1)
|
||||||
|
assert result == "Hello world"
|
||||||
|
|
||||||
|
async def test_respects_hard_limit(self, db: Database) -> None:
|
||||||
|
"""Output is truncated at hard_limit."""
|
||||||
|
words = [f"w{i}" for i in range(100)]
|
||||||
|
text = " ".join(words)
|
||||||
|
await ingest(db, 1, 100, text)
|
||||||
|
|
||||||
|
result = await generate(db, 1, soft_limit=10, hard_limit=49)
|
||||||
|
assert result == "w0 w1 w2 w3 w4 w5 w6 w7 w8 w9 w10 w11 w12 w13 w14"
|
||||||
|
|
||||||
|
async def test_prefers_completing_word_after_soft_limit(self, db: Database) -> None:
|
||||||
|
"""After soft_limit, generation prefers completing words."""
|
||||||
|
# Chain: A → B → C → D → {E, "end."}
|
||||||
|
# A, B, C have only non-completing transitions, so past the soft limit
|
||||||
|
# the loop falls back to get_random_next_word for each. D has both a
|
||||||
|
# continuing ("E") and completing ("end.") transition, so
|
||||||
|
# get_random_completing_next_word deterministically picks "end.".
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "A")],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "A", "B"),
|
||||||
|
Transition(1, 100, "B", "C"),
|
||||||
|
Transition(1, 100, "C", "D"),
|
||||||
|
Transition(1, 100, "D", "E"),
|
||||||
|
Transition(1, 100, "D", "F"),
|
||||||
|
Transition(1, 100, "D", "G"),
|
||||||
|
Transition(1, 100, "D", "H"),
|
||||||
|
Transition(1, 100, "D", "end."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
for _ in range(100):
|
||||||
|
result = await generate(db, 1, soft_limit=1, hard_limit=1000)
|
||||||
|
assert result == "A B C D end."
|
||||||
|
|
||||||
|
async def test_start_word_already_ends_sentence(self, db: Database) -> None:
|
||||||
|
"""Generation stops immediately when the start word is sentence-ending."""
|
||||||
|
await db.add_markov_data([StartWord(1, 100, "Yes.")], [])
|
||||||
|
|
||||||
|
result = await generate(db, 1)
|
||||||
|
assert result == "Yes."
|
||||||
|
|
||||||
|
async def test_chain_dead_end(self, db: Database) -> None:
|
||||||
|
"""Generation stops when no next word exists (dead-end chain)."""
|
||||||
|
await db.add_markov_data([StartWord(1, 100, "Hello")], [])
|
||||||
|
|
||||||
|
# "Hello" has no transitions, so the loop breaks immediately.
|
||||||
|
result = await generate(db, 1)
|
||||||
|
assert result == "Hello"
|
||||||
|
|
||||||
|
async def test_hard_limit_strips_section_sign(self, db: Database) -> None:
|
||||||
|
"""Section sign at the truncation boundary is stripped."""
|
||||||
|
# Build a chain: "A" -> "B" -> "C§".
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "A")],
|
||||||
|
[
|
||||||
|
Transition(1, 100, "A", "B"),
|
||||||
|
Transition(1, 100, "B", f"C{DEFAULT_SENTENCE_END}"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
# hard_limit=5 truncates "A B C§" (length 6) to "A B C".
|
||||||
|
result = await generate(db, 1, soft_limit=100, hard_limit=5)
|
||||||
|
assert result == "A B C"
|
||||||
|
|
||||||
|
async def test_soft_limit_falls_back_to_regular_next_word(
|
||||||
|
self,
|
||||||
|
db: Database,
|
||||||
|
) -> None:
|
||||||
|
"""After soft_limit, falls back when no completing word exists."""
|
||||||
|
# "A" -> "B" (no completing transition). Past soft_limit,
|
||||||
|
# get_random_completing_next_word returns None, falls back to "B".
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "A")],
|
||||||
|
[Transition(1, 100, "A", "B")],
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await generate(db, 1, soft_limit=1, hard_limit=1000)
|
||||||
|
assert result == "A B"
|
||||||
|
|
||||||
|
async def test_hard_limit_truncates_mid_word(self, db: Database) -> None:
|
||||||
|
"""Hard limit slices output even when it falls inside a word."""
|
||||||
|
await db.add_markov_data(
|
||||||
|
[StartWord(1, 100, "AB")],
|
||||||
|
[Transition(1, 100, "AB", "CDEF")],
|
||||||
|
)
|
||||||
|
|
||||||
|
# "AB CDEF" is 7 chars; hard_limit=5 truncates to "AB CD".
|
||||||
|
result = await generate(db, 1, soft_limit=100, hard_limit=5)
|
||||||
|
assert result == "AB CD"
|
||||||
|
|
||||||
|
async def test_channel_isolation(self, db: Database) -> None:
|
||||||
|
"""Data ingested into one channel does not leak into another."""
|
||||||
|
await ingest(db, 1, 100, "Channel one data.")
|
||||||
|
await ingest(db, 2, 100, "Channel two data.")
|
||||||
|
|
||||||
|
# Channel 3 has no data; should get the fallback.
|
||||||
|
result = await generate(db, 3)
|
||||||
|
expected = (
|
||||||
|
"I do not have enough data to generate a message yet."
|
||||||
|
" Chat a bit more so I can learn how this channel talks."
|
||||||
|
)
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
|
||||||
|
class TestUningestSentence:
|
||||||
|
"""Single sentence uningest from the Markov chain."""
|
||||||
|
|
||||||
|
async def test_removes_start_word_and_transition(self, db: Database) -> None:
|
||||||
|
"""Uningest removes the start word and transition added by ingest."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
await _uningest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
|
||||||
|
async def test_preserves_other_data(self, db: Database) -> None:
|
||||||
|
"""Uningest only removes data for the specified sentence."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
await _ingest_sentence(db, 1, 100, "Goodbye world.")
|
||||||
|
await _uningest_sentence(db, 1, 100, "Hello world.")
|
||||||
|
assert await db.get_random_start_word(1) == "Goodbye"
|
||||||
|
|
||||||
|
async def test_handles_missing_punctuation(self, db: Database) -> None:
|
||||||
|
"""Uningest appends the default sentence end, matching ingest behavior."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello world")
|
||||||
|
await _uningest_sentence(db, 1, 100, "Hello world")
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_single_word_is_noop(self, db: Database) -> None:
|
||||||
|
"""Uningesting a single-word sentence does not error."""
|
||||||
|
await _ingest_sentence(db, 1, 100, "Hello.")
|
||||||
|
await _uningest_sentence(db, 1, 100, "Hello.")
|
||||||
|
|
||||||
|
|
||||||
|
class TestUningestParagraph:
|
||||||
|
"""Paragraph-level uningest."""
|
||||||
|
|
||||||
|
async def test_uningest_multiple_sentences(self, db: Database) -> None:
|
||||||
|
"""Uningest reverses a multi-sentence paragraph."""
|
||||||
|
await ingest(db, 1, 100, "Hello world. Goodbye world!")
|
||||||
|
await uningest(db, 1, 100, "Hello world. Goodbye world!")
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_uningest_preserves_duplicate_data(self, db: Database) -> None:
|
||||||
|
"""Uningesting one copy leaves the other intact."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await uningest(db, 1, 100, "Hello world.")
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
@@ -0,0 +1,206 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for the Prometheus metrics module.
|
||||||
|
|
||||||
|
Tests cover metric object registration and the MetricsServer HTTP endpoint.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import errno
|
||||||
|
import socket
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from http import HTTPStatus
|
||||||
|
from typing import TYPE_CHECKING, Literal, cast
|
||||||
|
|
||||||
|
import aiohttp
|
||||||
|
import pytest
|
||||||
|
from prometheus_client import generate_latest
|
||||||
|
|
||||||
|
from crabstero.metrics import MetricsServer, TcpMetricsAddress, UnixMetricsAddress
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
type MetricsTransport = Literal["tcp", "unix-socket"]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class MetricsEndpoint:
|
||||||
|
"""Client details for one running metrics transport."""
|
||||||
|
|
||||||
|
base_url: str
|
||||||
|
unix_socket_path: str | None = None
|
||||||
|
|
||||||
|
def client_session(self) -> aiohttp.ClientSession:
|
||||||
|
"""Create an aiohttp client session for this metrics transport."""
|
||||||
|
connector = (
|
||||||
|
aiohttp.UnixConnector(path=self.unix_socket_path)
|
||||||
|
if self.unix_socket_path is not None
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
return aiohttp.ClientSession(connector=connector)
|
||||||
|
|
||||||
|
|
||||||
|
class TestMetricObjects:
|
||||||
|
"""All expected metric objects exist and are registered."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"name",
|
||||||
|
[
|
||||||
|
pytest.param("crabstero_build_info", id="build-info"),
|
||||||
|
pytest.param("crabstero_replies_sent_total", id="replies-sent"),
|
||||||
|
pytest.param("crabstero_embeds_generated_total", id="embeds-generated"),
|
||||||
|
pytest.param("crabstero_messages_processed_total", id="messages-processed"),
|
||||||
|
pytest.param("crabstero_pingme_total", id="pingme"),
|
||||||
|
pytest.param("crabstero_forgetme_total", id="forgetme"),
|
||||||
|
pytest.param("crabstero_errors_total", id="errors"),
|
||||||
|
pytest.param("crabstero_discord_latency_seconds", id="discord-latency"),
|
||||||
|
pytest.param("crabstero_guild_count", id="guild-count"),
|
||||||
|
pytest.param("crabstero_ingestion_active_channels", id="ingestion-active"),
|
||||||
|
pytest.param(
|
||||||
|
"crabstero_message_ingestion_duration_seconds",
|
||||||
|
id="message-ingestion-duration",
|
||||||
|
),
|
||||||
|
pytest.param(
|
||||||
|
"crabstero_channel_ingestion_duration_seconds",
|
||||||
|
id="channel-ingestion-duration",
|
||||||
|
),
|
||||||
|
pytest.param(
|
||||||
|
"crabstero_generation_duration_seconds",
|
||||||
|
id="generation-duration",
|
||||||
|
),
|
||||||
|
pytest.param(
|
||||||
|
"crabstero_channel_ingestion_messages",
|
||||||
|
id="channel-ingestion-messages",
|
||||||
|
),
|
||||||
|
pytest.param("crabstero_discord_events_total", id="discord-events"),
|
||||||
|
pytest.param(
|
||||||
|
"crabstero_messages_uningested_total",
|
||||||
|
id="messages-uningested",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_metric_in_output(self, name: str) -> None:
|
||||||
|
"""Each declared metric appears in the Prometheus text output."""
|
||||||
|
output = generate_latest().decode()
|
||||||
|
assert name in output, f"{name} not found in Prometheus output"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(
|
||||||
|
params=[
|
||||||
|
pytest.param("tcp", id="tcp"),
|
||||||
|
pytest.param("unix-socket", id="unix-socket"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def metrics_endpoint(
|
||||||
|
request: pytest.FixtureRequest,
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> AsyncGenerator[MetricsEndpoint]:
|
||||||
|
"""Start a MetricsServer on each supported transport."""
|
||||||
|
transport = cast("MetricsTransport", request.param)
|
||||||
|
match transport:
|
||||||
|
case "tcp":
|
||||||
|
server = MetricsServer(TcpMetricsAddress("127.0.0.1", 0))
|
||||||
|
await server.start()
|
||||||
|
endpoint = MetricsEndpoint(f"http://127.0.0.1:{server.port}")
|
||||||
|
case "unix-socket":
|
||||||
|
socket_path = tmp_path / "metrics.sock"
|
||||||
|
server = MetricsServer(UnixMetricsAddress(str(socket_path)))
|
||||||
|
await server.start()
|
||||||
|
endpoint = MetricsEndpoint(
|
||||||
|
"http://crabstero",
|
||||||
|
unix_socket_path=str(socket_path),
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
yield endpoint
|
||||||
|
finally:
|
||||||
|
await server.stop()
|
||||||
|
|
||||||
|
|
||||||
|
class TestMetricsServer:
|
||||||
|
"""HTTP server serves Prometheus metrics on /metrics."""
|
||||||
|
|
||||||
|
async def test_serves_metrics_endpoint(
|
||||||
|
self,
|
||||||
|
metrics_endpoint: MetricsEndpoint,
|
||||||
|
) -> None:
|
||||||
|
"""GET /metrics returns 200 with metric output containing our metrics."""
|
||||||
|
async with (
|
||||||
|
metrics_endpoint.client_session() as session,
|
||||||
|
session.get(f"{metrics_endpoint.base_url}/metrics") as resp,
|
||||||
|
):
|
||||||
|
assert resp.status == HTTPStatus.OK
|
||||||
|
body = await resp.text()
|
||||||
|
assert "crabstero_build_info" in body
|
||||||
|
|
||||||
|
async def test_non_metrics_path_returns_404(
|
||||||
|
self,
|
||||||
|
metrics_endpoint: MetricsEndpoint,
|
||||||
|
) -> None:
|
||||||
|
"""GET on an unknown path returns 404."""
|
||||||
|
async with (
|
||||||
|
metrics_endpoint.client_session() as session,
|
||||||
|
session.get(f"{metrics_endpoint.base_url}/notfound") as resp,
|
||||||
|
):
|
||||||
|
assert resp.status == HTTPStatus.NOT_FOUND
|
||||||
|
|
||||||
|
async def test_stale_unix_socket_path_is_recovered(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
"""Startup removes an abandoned Unix socket file left by a crash."""
|
||||||
|
socket_path = tmp_path / "metrics.sock"
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as stale_socket:
|
||||||
|
stale_socket.bind(str(socket_path))
|
||||||
|
|
||||||
|
server = MetricsServer(UnixMetricsAddress(str(socket_path)))
|
||||||
|
await server.start()
|
||||||
|
try:
|
||||||
|
connector = aiohttp.UnixConnector(path=str(socket_path))
|
||||||
|
async with (
|
||||||
|
aiohttp.ClientSession(connector=connector) as session,
|
||||||
|
session.get("http://crabstero/metrics") as resp,
|
||||||
|
):
|
||||||
|
assert resp.status == HTTPStatus.OK
|
||||||
|
body = await resp.text()
|
||||||
|
assert "crabstero_build_info" in body
|
||||||
|
finally:
|
||||||
|
await server.stop()
|
||||||
|
|
||||||
|
async def test_active_unix_socket_path_fails(
|
||||||
|
self,
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
"""Startup refuses to replace a Unix socket path that is still in use."""
|
||||||
|
socket_path = tmp_path / "metrics.sock"
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as active_socket:
|
||||||
|
active_socket.bind(str(socket_path))
|
||||||
|
active_socket.listen(1)
|
||||||
|
|
||||||
|
server = MetricsServer(UnixMetricsAddress(str(socket_path)))
|
||||||
|
with pytest.raises(OSError, match="already in use") as exc_info:
|
||||||
|
await server.start()
|
||||||
|
|
||||||
|
assert exc_info.value.errno == errno.EADDRINUSE
|
||||||
|
|
||||||
|
async def test_non_socket_unix_path_fails(self, tmp_path: Path) -> None:
|
||||||
|
"""Startup refuses to replace a non-socket path."""
|
||||||
|
socket_path = tmp_path / "metrics.sock"
|
||||||
|
socket_path.write_text("")
|
||||||
|
server = MetricsServer(UnixMetricsAddress(str(socket_path)))
|
||||||
|
with pytest.raises(FileExistsError):
|
||||||
|
await server.start()
|
||||||
@@ -0,0 +1,265 @@
|
|||||||
|
# Copyright 2026 Logan Fick
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Integration tests for the full ingest → uningest cycle."""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from crabstero.cache import CachedMessage, IngestCache
|
||||||
|
from crabstero.database import ChannelImage
|
||||||
|
from crabstero.markov import ingest, uningest
|
||||||
|
from crabstero.messages import uningest_message
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
import aiosqlite
|
||||||
|
|
||||||
|
from crabstero.database import Database
|
||||||
|
|
||||||
|
|
||||||
|
async def _snapshot(db: Database) -> dict[str, list[aiosqlite.Row]]:
|
||||||
|
"""Return sorted rows from all Markov-related tables."""
|
||||||
|
tables: dict[str, list[aiosqlite.Row]] = {}
|
||||||
|
for table in ("markov_start_words", "markov_transitions", "channel_images"):
|
||||||
|
async with db._connection.execute(f"SELECT * FROM {table}") as cursor: # noqa: S608
|
||||||
|
tables[table] = sorted(await cursor.fetchall())
|
||||||
|
return tables
|
||||||
|
|
||||||
|
|
||||||
|
class TestIngestUningestCycle:
|
||||||
|
"""Full round-trip: ingest data, then uningest it completely."""
|
||||||
|
|
||||||
|
async def test_content_round_trip(self, db: Database) -> None:
|
||||||
|
"""Ingest and uningest content leaves the database clean."""
|
||||||
|
await ingest(db, 1, 100, "Hello beautiful world.")
|
||||||
|
await uningest(db, 1, 100, "Hello beautiful world.")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
assert await db.get_random_next_word(1, "beautiful") is None
|
||||||
|
|
||||||
|
async def test_image_round_trip(self, db: Database) -> None:
|
||||||
|
"""Ingest and uningest an image leaves the database clean."""
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/cat.png")])
|
||||||
|
await db.remove_images([ChannelImage(1, 100, "https://example.com/cat.png")])
|
||||||
|
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
|
||||||
|
async def test_multi_sentence_round_trip(self, db: Database) -> None:
|
||||||
|
"""Multi-sentence ingest and uningest leaves the database clean."""
|
||||||
|
text = "Hello world. Goodbye world! How are you?"
|
||||||
|
await ingest(db, 1, 100, text)
|
||||||
|
await uningest(db, 1, 100, text)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_uningest_only_removes_one_copy(self, db: Database) -> None:
|
||||||
|
"""Uningesting once preserves data from a second identical ingest."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await uningest(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
# One copy of each row should remain.
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_uningest_does_not_affect_other_channels(self, db: Database) -> None:
|
||||||
|
"""Uningesting from one channel leaves another channel's data intact."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await ingest(db, 2, 100, "Hello world.")
|
||||||
|
await uningest(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_start_word(2) == "Hello"
|
||||||
|
assert await db.get_random_next_word(2, "Hello") == "world."
|
||||||
|
|
||||||
|
async def test_uningest_does_not_affect_other_users(self, db: Database) -> None:
|
||||||
|
"""Uningesting for one user leaves another user's data."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
await ingest(db, 1, 200, "Hello world.")
|
||||||
|
await uningest(db, 1, 100, "Hello world.")
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) == "Hello"
|
||||||
|
assert await db.get_random_next_word(1, "Hello") == "world."
|
||||||
|
|
||||||
|
|
||||||
|
class TestUningestRestoresState:
|
||||||
|
"""Uningest restores the database to its prior state."""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"text",
|
||||||
|
[
|
||||||
|
pytest.param("Hello world.", id="simple-sentence"),
|
||||||
|
pytest.param("Hello world", id="missing-punctuation"),
|
||||||
|
pytest.param(
|
||||||
|
"Hello world. Goodbye world! How are you?",
|
||||||
|
id="multi-sentence",
|
||||||
|
),
|
||||||
|
pytest.param("One.", id="single-word"),
|
||||||
|
pytest.param(
|
||||||
|
"Lots of extra spaces\nand\nnewlines here.",
|
||||||
|
id="whitespace-normalization",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_uningest_restores_empty_db(self, db: Database, text: str) -> None:
|
||||||
|
"""Ingest then uningest on an empty database leaves all tables empty."""
|
||||||
|
before = await _snapshot(db)
|
||||||
|
await ingest(db, 1, 100, text)
|
||||||
|
await uningest(db, 1, 100, text)
|
||||||
|
after = await _snapshot(db)
|
||||||
|
assert after == before
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"text",
|
||||||
|
[
|
||||||
|
pytest.param("Hello world.", id="simple-sentence"),
|
||||||
|
pytest.param(
|
||||||
|
"Hello world. Goodbye world! How are you?",
|
||||||
|
id="multi-sentence",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_uningest_restores_preexisting_data(
|
||||||
|
self,
|
||||||
|
db: Database,
|
||||||
|
text: str,
|
||||||
|
) -> None:
|
||||||
|
"""Ingest then uningest preserves unrelated pre-existing data exactly."""
|
||||||
|
await ingest(db, 99, 200, "Pre-existing data stays safe.")
|
||||||
|
await db.add_images([ChannelImage(99, 200, "https://example.com/existing.png")])
|
||||||
|
|
||||||
|
before = await _snapshot(db)
|
||||||
|
await ingest(db, 1, 100, text)
|
||||||
|
await uningest(db, 1, 100, text)
|
||||||
|
after = await _snapshot(db)
|
||||||
|
assert after == before
|
||||||
|
|
||||||
|
|
||||||
|
class TestUningestMessage:
|
||||||
|
"""Orchestrated uningest via cache lookup and database reversal."""
|
||||||
|
|
||||||
|
async def test_content_only(self, db: Database) -> None:
|
||||||
|
"""Uningest reverses a content-only message via the cache."""
|
||||||
|
await ingest(db, 1, 100, "Hello beautiful world.")
|
||||||
|
cache = IngestCache()
|
||||||
|
cache.put(
|
||||||
|
555,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello beautiful world.",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 555)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_next_word(1, "Hello") is None
|
||||||
|
|
||||||
|
async def test_embeds_only(self, db: Database) -> None:
|
||||||
|
"""Uningest reverses embed text ingestion."""
|
||||||
|
await ingest(db, 1, 100, "Embed title here.")
|
||||||
|
await ingest(db, 1, 100, "Embed description here.")
|
||||||
|
cache = IngestCache()
|
||||||
|
cache.put(
|
||||||
|
556,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content=None,
|
||||||
|
embed_texts=["Embed title here.", "Embed description here."],
|
||||||
|
image_urls=[],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 556)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
|
||||||
|
async def test_content_with_embeds_and_images(self, db: Database) -> None:
|
||||||
|
"""Uningest reverses content, embed text, and image data together."""
|
||||||
|
await ingest(db, 1, 100, "Body text here.")
|
||||||
|
await ingest(db, 1, 100, "Embed title.")
|
||||||
|
await db.add_images([ChannelImage(1, 100, "https://example.com/img.png")])
|
||||||
|
cache = IngestCache()
|
||||||
|
cache.put(
|
||||||
|
557,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Body text here.",
|
||||||
|
embed_texts=["Embed title."],
|
||||||
|
image_urls=["https://example.com/img.png"],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 557)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) is None
|
||||||
|
assert await db.get_random_image(1) is None
|
||||||
|
|
||||||
|
async def test_cache_miss_is_noop(self, db: Database) -> None:
|
||||||
|
"""A message not in the cache leaves the database unchanged."""
|
||||||
|
await ingest(db, 1, 100, "Keep this data.")
|
||||||
|
cache = IngestCache()
|
||||||
|
before = await _snapshot(db)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 999)
|
||||||
|
|
||||||
|
after = await _snapshot(db)
|
||||||
|
assert after == before
|
||||||
|
|
||||||
|
async def test_preserves_other_messages(self, db: Database) -> None:
|
||||||
|
"""Uningesting one message leaves another message's data intact."""
|
||||||
|
await ingest(db, 1, 100, "First message.")
|
||||||
|
await ingest(db, 1, 100, "Second message.")
|
||||||
|
cache = IngestCache()
|
||||||
|
cache.put(
|
||||||
|
601,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="First message.",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 601)
|
||||||
|
|
||||||
|
assert await db.get_random_start_word(1) == "Second"
|
||||||
|
assert await db.get_random_next_word(1, "Second") == "message."
|
||||||
|
|
||||||
|
async def test_pops_entry_from_cache(self, db: Database) -> None:
|
||||||
|
"""The cache entry is consumed after uningest."""
|
||||||
|
await ingest(db, 1, 100, "Hello world.")
|
||||||
|
cache = IngestCache()
|
||||||
|
cache.put(
|
||||||
|
602,
|
||||||
|
CachedMessage(
|
||||||
|
channel_id=1,
|
||||||
|
user_id=100,
|
||||||
|
content="Hello world.",
|
||||||
|
embed_texts=[],
|
||||||
|
image_urls=[],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
await uningest_message(db, cache, 602)
|
||||||
|
|
||||||
|
assert cache.pop(602) is None
|
||||||
@@ -12,9 +12,15 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unit tests for package version availability."""
|
||||||
|
|
||||||
import crabstero
|
import crabstero
|
||||||
|
|
||||||
|
|
||||||
def test_version_is_available() -> None:
|
class TestVersion:
|
||||||
|
"""Package version metadata."""
|
||||||
|
|
||||||
|
def test_version_is_available(self) -> None:
|
||||||
|
"""__version__ is a non-empty string."""
|
||||||
assert isinstance(crabstero.__version__, str)
|
assert isinstance(crabstero.__version__, str)
|
||||||
assert crabstero.__version__
|
assert crabstero.__version__
|
||||||
|
|||||||
Reference in New Issue
Block a user