Skip to main content

studio_worker/
cli.rs

1//! Clap CLI definitions, kept out of `main.rs` so they're testable.
2use clap::{Parser, Subcommand};
3
4#[derive(Parser, Debug)]
5#[command(
6    name = "studio-worker",
7    version,
8    about = "Studio worker — pull-based generation agent (image / llm / audio / video)"
9)]
10pub struct Cli {
11    /// Override the path to config.toml.
12    #[arg(long, global = true)]
13    pub config: Option<String>,
14    #[command(subcommand)]
15    pub command: Command,
16}
17
18#[derive(Subcommand, Debug, PartialEq)]
19pub enum Command {
20    /// Start the heartbeat + claim loop.
21    Run,
22    /// Register the worker against the API (idempotent).
23    Register {
24        #[arg(long)]
25        bootstrap_token: Option<String>,
26        #[arg(long)]
27        api_base_url: Option<String>,
28    },
29    /// Print local config + last heartbeat info.
30    Status,
31    /// Install platform-appropriate auto-start service.
32    InstallService,
33    /// Uninstall the auto-start service.
34    UninstallService,
35    /// Enable auto-claim.
36    Enable,
37    /// Disable auto-claim.
38    Disable,
39    /// Set the VRAM threshold (GB) the worker reports.
40    SetThreshold { gb: f32 },
41    /// Print resolved config + relevant paths.
42    Config,
43    /// Check the release feed for a newer version (does not install).
44    CheckUpdate,
45    /// Launch the desktop UI (requires the `ui` cargo feature).
46    Ui,
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use clap::Parser;
53
54    #[test]
55    fn parses_run() {
56        let cli = Cli::parse_from(["studio-worker", "run"]);
57        assert!(matches!(cli.command, Command::Run));
58        assert!(cli.config.is_none());
59    }
60
61    #[test]
62    fn parses_run_with_config_override() {
63        let cli = Cli::parse_from(["studio-worker", "--config", "/etc/x.toml", "run"]);
64        assert_eq!(cli.config.as_deref(), Some("/etc/x.toml"));
65        assert!(matches!(cli.command, Command::Run));
66    }
67
68    #[test]
69    fn parses_register_with_overrides() {
70        let cli = Cli::parse_from([
71            "studio-worker",
72            "register",
73            "--bootstrap-token",
74            "secret",
75            "--api-base-url",
76            "https://example.invalid",
77        ]);
78        match cli.command {
79            Command::Register {
80                bootstrap_token,
81                api_base_url,
82            } => {
83                assert_eq!(bootstrap_token.as_deref(), Some("secret"));
84                assert_eq!(api_base_url.as_deref(), Some("https://example.invalid"));
85            }
86            other => panic!("expected register, got {other:?}"),
87        }
88    }
89
90    #[test]
91    fn parses_set_threshold_with_float() {
92        let cli = Cli::parse_from(["studio-worker", "set-threshold", "12.5"]);
93        match cli.command {
94            Command::SetThreshold { gb } => assert!((gb - 12.5).abs() < 1e-6),
95            other => panic!("expected set-threshold, got {other:?}"),
96        }
97    }
98
99    #[test]
100    fn parses_all_simple_subcommands() {
101        let cases = [
102            ("status", Command::Status),
103            ("install-service", Command::InstallService),
104            ("uninstall-service", Command::UninstallService),
105            ("enable", Command::Enable),
106            ("disable", Command::Disable),
107            ("config", Command::Config),
108            ("check-update", Command::CheckUpdate),
109            ("ui", Command::Ui),
110        ];
111        for (name, expected) in cases {
112            let cli = Cli::parse_from(["studio-worker", name]);
113            assert_eq!(cli.command, expected, "parsing `{name}`");
114        }
115    }
116}