test(capabilities): cover registry, executor, schemas, billing

surfsense_backend/tests/unit/capabilities/test_billing.py

@@ -0,0 +1,116 @@
+"""Billing charges the workspace owner once per billable success at the executor (03c).
+
+Boundaries mocked: the DB session and the audit helper. NOT mocked: the real
+WebCrawlCreditService debit math and the owner-billed decision.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+from uuid import UUID
+
+import pytest
+
+import app.capabilities.billing as billing
+from app.capabilities.billing import charge_capability
+from app.capabilities.types import BillingUnit, CapabilityContext
+from app.capabilities.web.scrape.schemas import ScrapeOutput, ScrapeRow
+from app.config import config
+
+pytestmark = pytest.mark.unit
+
+_WORKSPACE_ID = 1
+_OWNER = UUID("00000000-0000-0000-0000-0000000000bb")
+
+
+class _FakeUser:
+    def __init__(self, balance_micros: int, reserved_micros: int = 0):
+        self.credit_micros_balance = balance_micros
+        self.credit_micros_reserved = reserved_micros
+
+
+def _make_session(owner_id, balance_micros):
+    """Mock session serving owner-resolution and the charge_credits debit."""
+    fake_user = _FakeUser(balance_micros)
+    session = AsyncMock()
+    session.add = MagicMock()
+
+    def _make_result(*_args, **_kwargs):
+        result = MagicMock()
+        result.scalar_one_or_none.return_value = owner_id  # owner resolution
+        result.unique.return_value.scalar_one_or_none.return_value = fake_user  # debit
+        return result
+
+    session.execute = AsyncMock(side_effect=_make_result)
+    return session, fake_user
+
+
+def _output(*statuses: str) -> ScrapeOutput:
+    return ScrapeOutput(
+        rows=[
+            ScrapeRow(url=f"https://{i}.com", status=status)
+            for i, status in enumerate(statuses)
+        ]
+    )
+
+
+def _ctx(session) -> CapabilityContext:
+    return CapabilityContext(session=session, workspace_id=_WORKSPACE_ID)
+
+
+@pytest.fixture(autouse=True)
+def _stub_auto_reload(monkeypatch):
+    import app.services.auto_reload_service as ar
+
+    monkeypatch.setattr(ar, "maybe_trigger_auto_reload", AsyncMock())
+
+
+@pytest.fixture
+def record_usage(monkeypatch):
+    rec = AsyncMock(return_value=MagicMock())
+    monkeypatch.setattr(billing, "record_token_usage", rec)
+    return rec
+
+
+async def test_charges_workspace_owner_per_successful_crawl(monkeypatch, record_usage):
+    monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
+    monkeypatch.setattr(config, "WEB_CRAWL_MICROS_PER_SUCCESS", 1000)
+    session, user = _make_session(_OWNER, balance_micros=100_000)
+
+    await charge_capability(
+        _output("success", "empty", "success"), BillingUnit.WEB_CRAWL, _ctx(session)
+    )
+
+    # Owner debited 2 * 1000; one web_crawl audit row billed to the OWNER.
+    assert user.credit_micros_balance == 100_000 - 2000
+    record_usage.assert_awaited_once()
+    kwargs = record_usage.await_args.kwargs
+    assert kwargs["usage_type"] == "web_crawl"
+    assert kwargs["user_id"] == _OWNER
+    assert kwargs["workspace_id"] == _WORKSPACE_ID
+    assert kwargs["cost_micros"] == 2000
+
+
+async def test_no_successful_rows_is_free(monkeypatch, record_usage):
+    monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", True)
+    session, user = _make_session(_OWNER, balance_micros=100_000)
+
+    await charge_capability(
+        _output("empty", "failed"), BillingUnit.WEB_CRAWL, _ctx(session)
+    )
+
+    record_usage.assert_not_awaited()
+    assert user.credit_micros_balance == 100_000
+
+
+async def test_disabled_is_noop(monkeypatch, record_usage):
+    monkeypatch.setattr(config, "WEB_CRAWL_CREDIT_BILLING_ENABLED", False)
+    session, user = _make_session(_OWNER, balance_micros=100_000)
+
+    await charge_capability(
+        _output("success", "success"), BillingUnit.WEB_CRAWL, _ctx(session)
+    )
+
+    record_usage.assert_not_awaited()
+    session.execute.assert_not_called()
+    assert user.credit_micros_balance == 100_000

surfsense_backend/tests/unit/capabilities/test_registry.py

@@ -0,0 +1,23 @@
+"""The registry exposes each verb as one Capability entry the doors/agent read from."""
+
+from __future__ import annotations
+
+import pytest
+
+from app.capabilities import (
+    web,  # noqa: F401  — importing the namespace registers its verbs
+)
+from app.capabilities.store import get_capability
+from app.capabilities.types import BillingUnit
+from app.capabilities.web.scrape.schemas import ScrapeInput, ScrapeOutput
+
+pytestmark = pytest.mark.unit
+
+
+def test_web_scrape_is_registered_with_its_schemas_and_billing_unit():
+    cap = get_capability("web.scrape")
+
+    assert cap.name == "web.scrape"
+    assert cap.input_schema is ScrapeInput
+    assert cap.output_schema is ScrapeOutput
+    assert cap.billing_unit is BillingUnit.WEB_CRAWL

surfsense_backend/tests/unit/capabilities/web/scrape/test_executor.py

@@ -0,0 +1,94 @@
+"""`web.scrape` executor behavior: URLs in → cleaned rows out.
+
+Boundary mocked: the crawler (injected fake). NOT mocked: the executor's own
+CrawlOutcome → ScrapeRow mapping.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from app.capabilities.web.scrape.executor import build_scrape_executor
+from app.capabilities.web.scrape.schemas import ScrapeInput, ScrapeOutput
+from app.proprietary.web_crawler import CrawlOutcome, CrawlOutcomeStatus
+
+pytestmark = pytest.mark.unit
+
+
+class _FakeCrawler:
+    """Stand-in for WebCrawlerConnector: serves a canned outcome per URL."""
+
+    def __init__(self, outcomes: dict[str, CrawlOutcome]):
+        self._outcomes = outcomes
+        self.calls: list[str] = []
+
+    async def crawl_url(self, url: str) -> CrawlOutcome:
+        self.calls.append(url)
+        return self._outcomes[url]
+
+
+def _success(content: str, metadata: dict[str, str]) -> CrawlOutcome:
+    return CrawlOutcome(
+        status=CrawlOutcomeStatus.SUCCESS,
+        result={
+            "content": content,
+            "metadata": metadata,
+            "crawler_type": "scrapling-static",
+        },
+    )
+
+
+async def test_scrape_returns_one_cleaned_row_for_a_successful_url():
+    url = "https://example.com"
+    crawler = _FakeCrawler({url: _success("# Hello", {"title": "Hello"})})
+    execute = build_scrape_executor(engine=crawler)
+
+    out = await execute(ScrapeInput(urls=[url]))
+
+    assert isinstance(out, ScrapeOutput)
+    assert len(out.rows) == 1
+    row = out.rows[0]
+    assert row.url == url
+    assert row.status == "success"
+    assert row.content == "# Hello"
+    assert row.metadata == {"title": "Hello"}
+
+
+async def test_scrape_returns_one_row_per_url_in_input_order():
+    a, b, c = "https://a.com", "https://b.com", "https://c.com"
+    crawler = _FakeCrawler(
+        {
+            a: _success("A", {"title": "A"}),
+            b: _success("B", {"title": "B"}),
+            c: _success("C", {"title": "C"}),
+        }
+    )
+    execute = build_scrape_executor(engine=crawler)
+
+    out = await execute(ScrapeInput(urls=[a, b, c]))
+
+    assert [row.url for row in out.rows] == [a, b, c]
+    assert [row.content for row in out.rows] == ["A", "B", "C"]
+
+
+async def test_partial_failure_keeps_the_batch_and_labels_each_url():
+    ok, empty, failed = "https://ok.com", "https://empty.com", "https://failed.com"
+    crawler = _FakeCrawler(
+        {
+            ok: _success("OK", {"title": "OK"}),
+            empty: CrawlOutcome(status=CrawlOutcomeStatus.EMPTY, error="no content"),
+            failed: CrawlOutcome(status=CrawlOutcomeStatus.FAILED, error="blocked"),
+        }
+    )
+    execute = build_scrape_executor(engine=crawler)
+
+    out = await execute(ScrapeInput(urls=[ok, empty, failed]))
+
+    by_url = {row.url: row for row in out.rows}
+    assert {u: r.status for u, r in by_url.items()} == {
+        ok: "success",
+        empty: "empty",
+        failed: "failed",
+    }
+    assert by_url[failed].content is None
+    assert by_url[failed].error == "blocked"

surfsense_backend/tests/unit/capabilities/web/scrape/test_schemas.py

@@ -0,0 +1,26 @@
+"""ScrapeOutput reports its own billable count (a success single-sources it)."""
+
+from __future__ import annotations
+
+import pytest
+
+from app.capabilities.web.scrape.schemas import ScrapeOutput, ScrapeRow
+
+pytestmark = pytest.mark.unit
+
+
+def _output(*statuses: str) -> ScrapeOutput:
+    return ScrapeOutput(
+        rows=[
+            ScrapeRow(url=f"https://{i}.com", status=status)
+            for i, status in enumerate(statuses)
+        ]
+    )
+
+
+def test_billable_units_counts_successful_rows():
+    assert _output("success", "empty", "success", "failed").billable_units == 2
+
+
+def test_billable_units_is_zero_without_successes():
+    assert _output("empty", "failed").billable_units == 0