feat: Implement unit specification and expansion logic in configuration

This commit is contained in:
2026-04-08 16:57:19 -07:00
parent 8be129259c
commit 01a09d6a61
4 changed files with 590 additions and 8 deletions

View File

@@ -27,4 +27,6 @@ from castle_core.manifest import ( # noqa: F401 — explicit re-exports for typ
ServiceSpec, ServiceSpec,
SystemdSpec, SystemdSpec,
TLSMode, TLSMode,
UnitKind,
UnitSpec,
) )

View File

@@ -4,12 +4,28 @@ from __future__ import annotations
import os import os
import re import re
from dataclasses import dataclass from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
import yaml import yaml
from castle_core.manifest import ProgramSpec, JobSpec, ServiceSpec from castle_core.manifest import (
CaddySpec,
DefaultsSpec,
ExposeSpec,
HttpExposeSpec,
HttpInternal,
JobSpec,
ManageSpec,
ProgramSpec,
ProxySpec,
RunCommand,
RunPython,
ServiceSpec,
SystemdSpec,
UnitKind,
UnitSpec,
)
CASTLE_HOME = Path.home() / ".castle" CASTLE_HOME = Path.home() / ".castle"
@@ -64,6 +80,8 @@ class CastleConfig:
programs: dict[str, ProgramSpec] programs: dict[str, ProgramSpec]
services: dict[str, ServiceSpec] services: dict[str, ServiceSpec]
jobs: dict[str, JobSpec] jobs: dict[str, JobSpec]
_unit_names: set[str] = field(default_factory=set)
_units_raw: dict[str, dict] = field(default_factory=dict)
@property @property
def tools(self) -> dict[str, ProgramSpec]: def tools(self) -> dict[str, ProgramSpec]:
@@ -129,6 +147,102 @@ def _parse_job(name: str, data: dict) -> JobSpec:
return JobSpec.model_validate(data_copy) return JobSpec.model_validate(data_copy)
def _parse_unit(name: str, data: dict) -> UnitSpec:
"""Parse a units: entry into a UnitSpec."""
data_copy = dict(data)
data_copy["id"] = name
return UnitSpec.model_validate(data_copy)
# Stack convention defaults used during unit expansion.
_STACK_DEFAULTS: dict[str, dict[str, str]] = {
"python-fastapi": {"runner": "python", "health_path": "/health"},
"python-cli": {"runner": "command"},
"react-vite": {},
}
# Kind → ProgramSpec.behavior mapping.
_KIND_BEHAVIOR: dict[str, str] = {
"tool": "tool",
"service": "daemon",
"site": "frontend",
"job": "tool",
}
def _expand_units(
units: dict[str, UnitSpec],
programs: dict[str, ProgramSpec],
services: dict[str, ServiceSpec],
jobs: dict[str, JobSpec],
) -> None:
"""Expand units: entries into programs/services/jobs dicts (in-place)."""
for name, unit in units.items():
if name in programs or name in services or name in jobs:
raise ValueError(
f"Unit '{name}' conflicts with existing entry in "
f"programs/services/jobs"
)
defaults = _STACK_DEFAULTS.get(unit.stack or "", {})
# Always create a ProgramSpec
programs[name] = ProgramSpec(
id=name,
description=unit.description,
behavior=_KIND_BEHAVIOR[unit.kind.value],
source=unit.source,
stack=unit.stack,
system_dependencies=unit.system_dependencies,
install_extras=unit.install_extras,
version=unit.version,
build=unit.build,
tags=unit.tags,
)
if unit.kind == UnitKind.SERVICE:
assert unit.port is not None # guaranteed by validator
runner = defaults.get("runner", "python")
run_spec: RunCommand | RunPython
if runner == "python":
run_spec = RunPython(runner="python", program=name)
else:
run_spec = RunCommand(runner="command", argv=[name])
health = unit.health_path or defaults.get("health_path")
services[name] = ServiceSpec(
id=name,
component=name,
run=run_spec,
expose=ExposeSpec(
http=HttpExposeSpec(
internal=HttpInternal(port=unit.port),
health_path=health,
)
),
proxy=ProxySpec(caddy=CaddySpec(path_prefix=unit.path_prefix))
if unit.path_prefix
else None,
manage=ManageSpec(systemd=SystemdSpec()),
defaults=DefaultsSpec(env=dict(unit.env)) if unit.env else None,
)
elif unit.kind == UnitKind.JOB:
assert unit.schedule is not None # guaranteed by validator
assert unit.argv is not None # guaranteed by validator
jobs[name] = JobSpec(
id=name,
component=name,
description=unit.description,
run=RunCommand(runner="command", argv=list(unit.argv)),
schedule=unit.schedule,
timezone=unit.timezone,
manage=ManageSpec(systemd=SystemdSpec()),
defaults=DefaultsSpec(env=dict(unit.env)) if unit.env else None,
)
def load_config(root: Path | None = None) -> CastleConfig: def load_config(root: Path | None = None) -> CastleConfig:
"""Load castle.yaml and return parsed configuration.""" """Load castle.yaml and return parsed configuration."""
if root is None: if root is None:
@@ -171,6 +285,25 @@ def load_config(root: Path | None = None) -> CastleConfig:
for name, job_data in data.get("jobs", {}).items(): for name, job_data in data.get("jobs", {}).items():
jobs[name] = _parse_job(name, job_data) jobs[name] = _parse_job(name, job_data)
# Expand units: section into programs/services/jobs
units_data = data.get("units") or {}
units: dict[str, UnitSpec] = {}
for name, unit_data in units_data.items():
units[name] = _parse_unit(name, unit_data)
unit_names: set[str] = set()
if units:
_expand_units(units, programs, services, jobs)
unit_names = set(units.keys())
# Resolve source paths for unit-generated programs
for name in unit_names:
prog = programs[name]
if prog.source:
if prog.source.startswith("repo:") and repo_path:
prog.source = str(repo_path / prog.source[5:])
elif not Path(prog.source).is_absolute():
prog.source = str(root / prog.source)
return CastleConfig( return CastleConfig(
root=root, root=root,
repo=repo_path, repo=repo_path,
@@ -178,6 +311,8 @@ def load_config(root: Path | None = None) -> CastleConfig:
programs=programs, programs=programs,
services=services, services=services,
jobs=jobs, jobs=jobs,
_unit_names=unit_names,
_units_raw=dict(units_data),
) )
@@ -258,9 +393,17 @@ def save_config(config: CastleConfig) -> None:
if config.repo: if config.repo:
data["repo"] = str(config.repo) data["repo"] = str(config.repo)
if config.programs: # Write units: section (raw roundtrip preserves user's original YAML)
if config._units_raw:
data["units"] = dict(config._units_raw)
# Write programs: (excluding unit-expanded ones)
non_unit_programs = {
k: v for k, v in config.programs.items() if k not in config._unit_names
}
if non_unit_programs:
data["programs"] = {} data["programs"] = {}
for name, spec in config.programs.items(): for name, spec in non_unit_programs.items():
d = _spec_to_yaml_dict(spec) d = _spec_to_yaml_dict(spec)
# Store relative source paths in YAML # Store relative source paths in YAML
if d.get("source") and Path(d["source"]).is_absolute(): if d.get("source") and Path(d["source"]).is_absolute():
@@ -280,14 +423,22 @@ def save_config(config: CastleConfig) -> None:
pass # not under root — keep absolute pass # not under root — keep absolute
data["programs"][name] = d data["programs"][name] = d
if config.services: # Write services: (excluding unit-expanded ones)
non_unit_services = {
k: v for k, v in config.services.items() if k not in config._unit_names
}
if non_unit_services:
data["services"] = {} data["services"] = {}
for name, spec in config.services.items(): for name, spec in non_unit_services.items():
data["services"][name] = _spec_to_yaml_dict(spec) data["services"][name] = _spec_to_yaml_dict(spec)
if config.jobs: # Write jobs: (excluding unit-expanded ones)
non_unit_jobs = {
k: v for k, v in config.jobs.items() if k not in config._unit_names
}
if non_unit_jobs:
data["jobs"] = {} data["jobs"] = {}
for name, spec in config.jobs.items(): for name, spec in non_unit_jobs.items():
data["jobs"][name] = _spec_to_yaml_dict(spec) data["jobs"][name] = _spec_to_yaml_dict(spec)
config_path = config.root / "castle.yaml" config_path = config.root / "castle.yaml"

View File

@@ -247,3 +247,58 @@ class JobSpec(BaseModel):
manage: ManageSpec | None = None manage: ManageSpec | None = None
defaults: DefaultsSpec | None = None defaults: DefaultsSpec | None = None
# ---------------------
# Unit spec — unified definition
# ---------------------
class UnitKind(str, Enum):
TOOL = "tool"
SERVICE = "service"
SITE = "site"
JOB = "job"
class UnitSpec(BaseModel):
"""Unified program+deployment definition.
Expands into ProgramSpec + optional ServiceSpec/JobSpec at config load time.
"""
id: str = ""
kind: UnitKind
description: str | None = None
# Program identity
source: str | None = None
stack: str | None = None
system_dependencies: list[str] = Field(default_factory=list)
install_extras: list[str] = Field(default_factory=list)
version: str | None = None
build: BuildSpec | None = None
tags: list[str] = Field(default_factory=list)
# Service fields (kind=service)
port: int | None = Field(default=None, ge=1, le=65535)
path_prefix: str | None = None
health_path: str | None = None
# Job fields (kind=job)
schedule: str | None = None
timezone: str = "America/Los_Angeles"
argv: list[str] | None = None
# Simplified env (flat dict, replaces defaults.env nesting)
env: EnvMap = Field(default_factory=dict)
@model_validator(mode="after")
def _validate_kind_fields(self) -> UnitSpec:
if self.kind == UnitKind.SERVICE and self.port is None:
raise ValueError("kind=service requires 'port'")
if self.kind == UnitKind.JOB and self.schedule is None:
raise ValueError("kind=job requires 'schedule'")
if self.kind == UnitKind.JOB and self.argv is None:
raise ValueError("kind=job requires 'argv'")
return self

374
core/tests/test_units.py Normal file
View File

@@ -0,0 +1,374 @@
"""Tests for units: expansion."""
from __future__ import annotations
from pathlib import Path
import pytest
import yaml
from castle_core.config import load_config, save_config
from castle_core.manifest import UnitKind, UnitSpec
@pytest.fixture
def units_root(tmp_path: Path) -> Path:
"""Castle root with a units: section."""
config = {
"gateway": {"port": 18000},
"units": {
"my-svc": {
"kind": "service",
"stack": "python-fastapi",
"description": "Test service",
"source": "code/my-svc",
"port": 9050,
"path_prefix": "/my-svc",
},
"my-tool": {
"kind": "tool",
"stack": "python-cli",
"description": "Test tool",
"source": "code/my-tool",
"system_dependencies": ["pandoc"],
},
"my-site": {
"kind": "site",
"stack": "react-vite",
"description": "Test frontend",
"source": "code/my-site",
"build": {
"commands": [["pnpm", "build"]],
"outputs": ["dist/"],
},
},
"my-job": {
"kind": "job",
"stack": "python-cli",
"description": "Test job",
"source": "code/my-job",
"schedule": "0 2 * * *",
"argv": ["my-job", "run"],
"env": {"DATA_DIR": "/tmp/data"},
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config, default_flow_style=False)
)
return tmp_path
class TestUnitExpansion:
def test_service_creates_program_and_service(self, units_root: Path) -> None:
config = load_config(units_root)
assert "my-svc" in config.programs
assert "my-svc" in config.services
prog = config.programs["my-svc"]
assert prog.behavior == "daemon"
assert prog.stack == "python-fastapi"
assert prog.description == "Test service"
svc = config.services["my-svc"]
assert svc.run.runner == "python"
assert svc.run.program == "my-svc"
assert svc.expose.http.internal.port == 9050
assert svc.expose.http.health_path == "/health"
assert svc.proxy.caddy.path_prefix == "/my-svc"
assert svc.manage.systemd is not None
assert svc.component == "my-svc"
def test_tool_creates_program_only(self, units_root: Path) -> None:
config = load_config(units_root)
assert "my-tool" in config.programs
assert "my-tool" not in config.services
assert "my-tool" not in config.jobs
prog = config.programs["my-tool"]
assert prog.behavior == "tool"
assert prog.system_dependencies == ["pandoc"]
def test_site_creates_program_with_build(self, units_root: Path) -> None:
config = load_config(units_root)
assert "my-site" in config.programs
assert "my-site" not in config.services
prog = config.programs["my-site"]
assert prog.behavior == "frontend"
assert prog.build is not None
assert prog.build.outputs == ["dist/"]
def test_job_creates_program_and_job(self, units_root: Path) -> None:
config = load_config(units_root)
assert "my-job" in config.programs
assert "my-job" in config.jobs
assert "my-job" not in config.services
prog = config.programs["my-job"]
assert prog.behavior == "tool"
job = config.jobs["my-job"]
assert job.schedule == "0 2 * * *"
assert job.run.runner == "command"
assert job.run.argv == ["my-job", "run"]
assert job.defaults.env["DATA_DIR"] == "/tmp/data"
assert job.component == "my-job"
def test_service_without_path_prefix(self, tmp_path: Path) -> None:
config_data = {
"gateway": {"port": 18000},
"units": {
"internal-svc": {
"kind": "service",
"stack": "python-fastapi",
"port": 9060,
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
svc = config.services["internal-svc"]
assert svc.proxy is None
def test_service_with_env(self, tmp_path: Path) -> None:
config_data = {
"gateway": {"port": 18000},
"units": {
"svc-env": {
"kind": "service",
"stack": "python-fastapi",
"port": 9060,
"env": {"API_KEY": "test123"},
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
svc = config.services["svc-env"]
assert svc.defaults.env["API_KEY"] == "test123"
def test_service_custom_health_path(self, tmp_path: Path) -> None:
config_data = {
"gateway": {"port": 18000},
"units": {
"custom-health": {
"kind": "service",
"stack": "python-fastapi",
"port": 9060,
"health_path": "/status",
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
svc = config.services["custom-health"]
assert svc.expose.http.health_path == "/status"
def test_source_paths_resolved(self, units_root: Path) -> None:
config = load_config(units_root)
prog = config.programs["my-tool"]
assert prog.source == str(units_root / "code" / "my-tool")
def test_repo_source_paths_resolved(self, tmp_path: Path) -> None:
repo_dir = tmp_path / "repo"
repo_dir.mkdir()
config_data = {
"gateway": {"port": 18000},
"repo": str(repo_dir),
"units": {
"repo-svc": {
"kind": "service",
"stack": "python-fastapi",
"port": 9060,
"source": "repo:my-api",
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
assert config.programs["repo-svc"].source == str(repo_dir / "my-api")
class TestUnitValidation:
def test_service_requires_port(self) -> None:
with pytest.raises(ValueError, match="requires 'port'"):
UnitSpec(id="bad", kind=UnitKind.SERVICE, stack="python-fastapi")
def test_job_requires_schedule(self) -> None:
with pytest.raises(ValueError, match="requires 'schedule'"):
UnitSpec(id="bad", kind=UnitKind.JOB, argv=["x"])
def test_job_requires_argv(self) -> None:
with pytest.raises(ValueError, match="requires 'argv'"):
UnitSpec(id="bad", kind=UnitKind.JOB, schedule="0 * * * *")
def test_tool_needs_no_port(self) -> None:
unit = UnitSpec(id="ok", kind=UnitKind.TOOL, stack="python-cli")
assert unit.port is None
def test_site_needs_no_port(self) -> None:
unit = UnitSpec(id="ok", kind=UnitKind.SITE, stack="react-vite")
assert unit.port is None
class TestUnitConflicts:
def test_conflict_with_existing_program(self, tmp_path: Path) -> None:
config = {
"gateway": {"port": 18000},
"programs": {
"conflict": {
"description": "Existing",
"behavior": "tool",
},
},
"units": {
"conflict": {
"kind": "tool",
"stack": "python-cli",
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config, default_flow_style=False)
)
with pytest.raises(ValueError, match="conflicts"):
load_config(tmp_path)
def test_conflict_with_existing_service(self, tmp_path: Path) -> None:
config = {
"gateway": {"port": 18000},
"services": {
"conflict": {
"run": {"runner": "command", "argv": ["svc"]},
},
},
"units": {
"conflict": {
"kind": "service",
"stack": "python-fastapi",
"port": 9999,
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config, default_flow_style=False)
)
with pytest.raises(ValueError, match="conflicts"):
load_config(tmp_path)
class TestUnitCoexistence:
def test_units_and_explicit_coexist(self, tmp_path: Path) -> None:
config = {
"gateway": {"port": 18000},
"programs": {
"manual-tool": {
"description": "Manual",
"behavior": "tool",
},
},
"services": {
"manual-svc": {
"run": {"runner": "command", "argv": ["svc"]},
"expose": {"http": {"internal": {"port": 19000}}},
},
},
"units": {
"unit-svc": {
"kind": "service",
"stack": "python-fastapi",
"port": 9070,
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config, default_flow_style=False)
)
cfg = load_config(tmp_path)
assert "manual-tool" in cfg.programs
assert "manual-svc" in cfg.services
assert "unit-svc" in cfg.programs
assert "unit-svc" in cfg.services
class TestUnitRoundTrip:
def test_save_preserves_units_section(self, units_root: Path) -> None:
config = load_config(units_root)
save_config(config)
config2 = load_config(units_root)
assert "my-svc" in config2.programs
assert "my-svc" in config2.services
assert "my-tool" in config2.programs
assert "my-job" in config2.jobs
def test_save_does_not_duplicate_into_explicit_sections(
self, units_root: Path
) -> None:
config = load_config(units_root)
save_config(config)
raw = yaml.safe_load((units_root / "castle.yaml").read_text())
# Unit entries should NOT appear in programs/services/jobs sections
assert "my-svc" not in raw.get("programs", {})
assert "my-svc" not in raw.get("services", {})
assert "my-job" not in raw.get("jobs", {})
# They should still be in units:
assert "my-svc" in raw.get("units", {})
assert "my-tool" in raw.get("units", {})
assert "my-job" in raw.get("units", {})
def test_mixed_save_separates_correctly(self, tmp_path: Path) -> None:
config_data = {
"gateway": {"port": 18000},
"programs": {
"manual-tool": {
"description": "Manual",
"behavior": "tool",
"source": "code/manual",
},
},
"units": {
"unit-tool": {
"kind": "tool",
"stack": "python-cli",
"source": "code/unit-tool",
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
save_config(config)
raw = yaml.safe_load((tmp_path / "castle.yaml").read_text())
# manual-tool stays in programs, unit-tool stays in units
assert "manual-tool" in raw.get("programs", {})
assert "unit-tool" not in raw.get("programs", {})
assert "unit-tool" in raw.get("units", {})
class TestNoUnitsSection:
def test_config_without_units_works(self, tmp_path: Path) -> None:
config_data = {
"gateway": {"port": 18000},
"programs": {
"some-tool": {
"description": "A tool",
"behavior": "tool",
},
},
}
(tmp_path / "castle.yaml").write_text(
yaml.dump(config_data, default_flow_style=False)
)
config = load_config(tmp_path)
assert "some-tool" in config.programs
assert not config._unit_names