@agentikos/omega-os 0.1.0 → 0.19.5

package/omega/Agentik_Engine/tests/test_account.py

@@ -0,0 +1,338 @@
+"""The account pool + vault + billing aggregator — the real things.
+
+End-to-end coverage in stdlib:
+  - load a YAML pool from a tempdir-based OMEGA_HOME
+  - round-robin and least-used selection
+  - read_token / write_token round-trip via the vault (mode 600 enforced)
+  - BillingAggregator aggregates per-account totals from seeded events
+  - the new providers (GLM, OpenAI, DeepSeek) import cleanly with no API key
+    and only raise on `.run()` — with a helpful message
+
+Standalone:  python3 tests/test_account.py
+"""
+import os
+import stat
+import sys
+import tempfile
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from omega_engine.account import (             # noqa: E402
+    AccountPool,
+    BillingAggregator,
+    ClaudeAccount,
+    read_token,
+    vault_path,
+    write_token,
+)
+from omega_engine.events import Event, EventType  # noqa: E402
+from omega_engine.provider import (             # noqa: E402
+    AgentRequest,
+    DeepSeekProvider,
+    GLMProvider,
+    OpenAIProvider,
+)
+from omega_engine.store import SQLiteStore     # noqa: E402
+
+
+# ──── helpers ────
+
+
+def _omega_home_with_pool(tmp: Path, selection: str = "least-used") -> Path:
+    """Seed OMEGA_HOME with an accounts.yaml containing two entries."""
+    cfg_dir = tmp / "Agentik_Providers" / "claude"
+    cfg_dir.mkdir(parents=True)
+    (cfg_dir / "accounts.yaml").write_text(
+        f"""version: 1
+selection: {selection}
+pool:
+  - id: max-primary
+    label: "Primary"
+    secret_ref: CLAUDE_OAUTH_max-primary
+    weight: 1
+    status: active
+  - id: max-secondary
+    label: "Secondary"
+    secret_ref: CLAUDE_OAUTH_max-secondary
+    weight: 2
+    status: active
+"""
+    )
+    return tmp
+
+
+# ──── tests ────
+
+
+def test_pool_load_yaml_two_accounts():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp))
+        pool = AccountPool.load(home)
+        assert len(pool.accounts) == 2, f"want 2, got {len(pool.accounts)}"
+        assert {a.id for a in pool.accounts} == {"max-primary", "max-secondary"}
+        assert pool.get("max-primary").label == "Primary"
+        assert pool.get("max-secondary").weight == 2
+
+
+def test_pool_falls_back_to_example_when_no_yaml():
+    # an OMEGA_HOME with only the example file should still load (the example
+    # is the template the repo ships with).
+    with tempfile.TemporaryDirectory() as tmp:
+        cfg = Path(tmp) / "Agentik_Providers" / "claude"
+        cfg.mkdir(parents=True)
+        (cfg / "accounts.example.yaml").write_text(
+            "version: 1\nselection: least-used\npool:\n  - id: max-x\n    "
+            "label: example\n    secret_ref: CLAUDE_OAUTH_max-x\n    "
+            "weight: 1\n    status: active\n"
+        )
+        pool = AccountPool.load(tmp)
+        assert len(pool.accounts) == 1
+        assert pool.accounts[0].id == "max-x"
+
+
+def test_pool_empty_when_no_config():
+    with tempfile.TemporaryDirectory() as tmp:
+        pool = AccountPool.load(tmp)
+        assert pool.accounts == []
+        # next() on an empty pool must raise — never silently return nothing
+        try:
+            pool.next()
+        except RuntimeError as exc:
+            assert "no active" in str(exc)
+        else:
+            raise AssertionError("expected RuntimeError on empty pool")
+
+
+def test_round_robin_alternates():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="round-robin")
+        pool = AccountPool.load(home)
+        picks = [pool.next().id for _ in range(4)]
+        # strict alternation regardless of usage counters
+        assert picks == [
+            "max-primary", "max-secondary", "max-primary", "max-secondary",
+        ], picks
+
+
+def test_least_used_picks_smaller_counter():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="least-used")
+        pool = AccountPool.load(home)
+        # bias primary to have more usage; secondary must win.
+        pool.usage_for("max-primary", 1000)
+        pick = pool.next()
+        assert pick.id == "max-secondary"
+        # ties: both at the same count → either is acceptable, but
+        # consistent picks one and the next call goes to the other.
+        pool.usage_for("max-secondary", 1000)
+        first = pool.next().id
+        second = pool.next().id
+        assert {first, second} <= {"max-primary", "max-secondary"}
+
+
+def test_set_status_rejects_bad_value():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp))
+        pool = AccountPool.load(home)
+        try:
+            pool.set_status("max-primary", "frozen")
+        except ValueError as exc:
+            assert "active" in str(exc)
+        else:
+            raise AssertionError("expected ValueError on bad status")
+        # but legal values work and persist through save/load
+        pool.set_status("max-primary", "resting")
+        pool.save(home)
+        reloaded = AccountPool.load(home)
+        assert reloaded.get("max-primary").status == "resting"
+
+
+def test_disabled_account_excluded_from_rotation():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="round-robin")
+        pool = AccountPool.load(home)
+        pool.set_status("max-secondary", "disabled")
+        picks = [pool.next().id for _ in range(3)]
+        assert all(p == "max-primary" for p in picks), picks
+
+
+def test_vault_round_trip_with_mode_600():
+    with tempfile.TemporaryDirectory() as tmp:
+        secret = "CLAUDE_OAUTH_max-primary"
+        path = write_token(tmp, secret, "sk-test-12345")
+        # exact file mode must be 0o600 — nothing else should be readable
+        mode = stat.S_IMODE(path.stat().st_mode)
+        assert mode == 0o600, f"want 0o600, got {oct(mode)}"
+        # The new vault writes to "<secret>_TOKEN[.age]"; the legacy
+        # vault_path("<secret>") still resolves to the dotenv path for
+        # backward compat reads — they may differ, but both must live under
+        # the same secrets dir.
+        secrets_dir = Path(tmp) / "Agentik_Extra" / "etc" / "secrets"
+        assert path.parent == secrets_dir, (path.parent, secrets_dir)
+        assert vault_path(tmp, secret).parent == secrets_dir
+        # round-trip — read_token tries the new vault first, then the legacy file
+        token = read_token(tmp, secret)
+        assert token == "sk-test-12345"
+
+
+def test_vault_refuses_path_traversal():
+    with tempfile.TemporaryDirectory() as tmp:
+        for bad in ("../escape", "a/b", "..", ""):
+            try:
+                vault_path(tmp, bad)
+            except ValueError:
+                continue
+            raise AssertionError(f"vault_path accepted bad ref: {bad!r}")
+
+
+def test_vault_read_missing_raises():
+    with tempfile.TemporaryDirectory() as tmp:
+        try:
+            read_token(tmp, "CLAUDE_OAUTH_nonexistent")
+        except FileNotFoundError:
+            return
+        raise AssertionError("expected FileNotFoundError on missing vault file")
+
+
+def test_billing_aggregates_per_account_from_events():
+    """Seed a real SQLite event store with task.* events that carry
+    `usage.account_id`, then verify BillingAggregator sums them correctly."""
+    with tempfile.TemporaryDirectory() as tmp:
+        db = Path(tmp) / "omega.db"
+        store = SQLiteStore(db)
+        # Two accounts billed across three task events. One event with no
+        # usage payload must be ignored (negative test in the same pass).
+        store.append(Event(
+            task_id="t-1", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-primary",
+                                "input_tokens": 100, "output_tokens": 200}},
+        ))
+        store.append(Event(
+            task_id="t-2", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-primary",
+                                "input_tokens": 50, "output_tokens": 75}},
+        ))
+        store.append(Event(
+            task_id="t-3", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-secondary",
+                                "input_tokens": 10, "output_tokens": 20}},
+        ))
+        # noise: an event with no usage payload — must NOT be counted
+        store.append(Event(
+            task_id="t-4", type=EventType.HEARTBEAT, payload={},
+        ))
+        # noise: an event with usage but zero tokens — must NOT be counted
+        store.append(Event(
+            task_id="t-5", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-secondary",
+                                "input_tokens": 0, "output_tokens": 0}},
+        ))
+        totals = BillingAggregator.from_event_log(store)
+        assert set(totals.keys()) == {"max-primary", "max-secondary"}, totals
+        assert totals["max-primary"] == {
+            "input_tokens": 150, "output_tokens": 275, "total_calls": 2,
+        }
+        assert totals["max-secondary"] == {
+            "input_tokens": 10, "output_tokens": 20, "total_calls": 1,
+        }
+        store.close()
+
+
+def test_billing_empty_log_returns_empty_dict():
+    with tempfile.TemporaryDirectory() as tmp:
+        db = Path(tmp) / "omega.db"
+        store = SQLiteStore(db)
+        totals = BillingAggregator.from_event_log(store)
+        assert totals == {}
+        store.close()
+
+
+def test_new_providers_import_without_api_key():
+    """The three new providers must be constructible with no env, no args.
+    Errors are deferred to `.run()` so the rest of the engine can keep
+    talking about them as objects."""
+    # cover the no-arg path explicitly
+    for cls in (GLMProvider, OpenAIProvider, DeepSeekProvider):
+        p = cls()
+        assert isinstance(p.id, str) and p.id
+        # default model is set
+        assert p._model, f"{cls.__name__} default model not set"
+        # custom args also accepted
+        p2 = cls(model="custom", api_key=None, base_url="https://example.test")
+        assert p2._model == "custom"
+        assert p2._base_url == "https://example.test"
+
+
+def _expect_runtime_error_on_run(provider, env_var: str) -> None:
+    """Calling .run() with no key in env AND no api_key arg must raise
+    a RuntimeError whose message names the missing env var."""
+    saved = os.environ.pop(env_var, None)
+    try:
+        provider.run(AgentRequest(role="worker", prompt="hi"))
+    except RuntimeError as exc:
+        msg = str(exc)
+        assert env_var in msg, f"{env_var} not in error message: {msg!r}"
+        return
+    finally:
+        if saved is not None:
+            os.environ[env_var] = saved
+    raise AssertionError(
+        f"{type(provider).__name__}.run() did not raise on missing key"
+    )
+
+
+def test_glm_provider_raises_without_key():
+    _expect_runtime_error_on_run(GLMProvider(), "GLM_API_KEY")
+
+
+def test_openai_provider_raises_without_key():
+    _expect_runtime_error_on_run(OpenAIProvider(), "OPENAI_API_KEY")
+
+
+def test_deepseek_provider_raises_without_key():
+    _expect_runtime_error_on_run(DeepSeekProvider(), "DEEPSEEK_API_KEY")
+
+
+def test_account_dataclass_round_trip_yaml():
+    """ClaudeAccount.to_yaml -> AccountPool.save -> AccountPool.load must
+    preserve all the fields we serialize."""
+    with tempfile.TemporaryDirectory() as tmp:
+        pool = AccountPool(
+            accounts=[
+                ClaudeAccount(
+                    id="max-new", label="round trip", secret_ref="CLAUDE_OAUTH_max-new",
+                    weight=3, status="resting", tokens_used=1234, last_used_at=1700.0,
+                ),
+            ],
+            selection="round-robin",
+        )
+        cfg = pool.save(tmp)
+        assert cfg.exists()
+        reloaded = AccountPool.load(tmp)
+        assert reloaded.selection == "round-robin"
+        a = reloaded.get("max-new")
+        assert a is not None
+        assert a.weight == 3
+        assert a.status == "resting"
+        assert a.tokens_used == 1234
+        assert a.last_used_at == 1700.0
+
+
+# ──── runner ────
+
+
+def _run_all() -> bool:
+    tests = [v for k, v in sorted(globals().items())
+             if k.startswith("test_") and callable(v)]
+    passed = 0
+    for t in tests:
+        t()
+        print(f"  PASS  {t.__name__}")
+        passed += 1
+    print(f"\n{passed}/{len(tests)} account tests passed")
+    return passed == len(tests)
+
+
+if __name__ == "__main__":
+    sys.exit(0 if _run_all() else 1)