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    /// Pre-set registration metadata before the next launch.
23    ///
24    /// On a fresh install you don't need this — `run` and `ui`
25    /// auto-register themselves.  Use it explicitly to:
26    ///   * set the human label shown in the studio's Pending
27    ///     Workers panel (`--label`)
28    ///   * point the worker at a different studio (`--api-base-url`)
29    ///   * clear local registration state after a rejection or
30    ///     between studios (`--reset`)
31    Register {
32        #[arg(long)]
33        api_base_url: Option<String>,
34        #[arg(long)]
35        label: Option<String>,
36        #[arg(long)]
37        reset: bool,
38    },
39    /// Print local config + last heartbeat info.
40    Status,
41    /// Install platform-appropriate auto-start service.
42    InstallService,
43    /// Uninstall the auto-start service.
44    UninstallService,
45    /// Enable auto-claim.
46    Enable,
47    /// Disable auto-claim.
48    Disable,
49    /// Set the VRAM threshold (GB) the worker reports.
50    SetThreshold { gb: f32 },
51    /// Print resolved config + relevant paths.
52    Config,
53    /// Check the release feed for a newer version (does not install).
54    CheckUpdate,
55    /// Launch the desktop UI (requires the `ui` cargo feature).
56    Ui,
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use clap::Parser;
63
64    #[test]
65    fn parses_run() {
66        let cli = Cli::parse_from(["studio-worker", "run"]);
67        assert!(matches!(cli.command, Command::Run));
68        assert!(cli.config.is_none());
69    }
70
71    #[test]
72    fn parses_run_with_config_override() {
73        let cli = Cli::parse_from(["studio-worker", "--config", "/etc/x.toml", "run"]);
74        assert_eq!(cli.config.as_deref(), Some("/etc/x.toml"));
75        assert!(matches!(cli.command, Command::Run));
76    }
77
78    #[test]
79    fn parses_register_with_overrides() {
80        let cli = Cli::parse_from([
81            "studio-worker",
82            "register",
83            "--api-base-url",
84            "https://example.invalid",
85        ]);
86        match cli.command {
87            Command::Register {
88                api_base_url,
89                label,
90                reset,
91            } => {
92                assert_eq!(api_base_url.as_deref(), Some("https://example.invalid"));
93                assert!(label.is_none());
94                assert!(!reset);
95            }
96            other => panic!("expected register, got {other:?}"),
97        }
98    }
99
100    #[test]
101    fn parses_register_with_label_and_reset() {
102        let cli = Cli::parse_from([
103            "studio-worker",
104            "register",
105            "--label",
106            "alice's rig",
107            "--reset",
108        ]);
109        match cli.command {
110            Command::Register {
111                api_base_url,
112                label,
113                reset,
114            } => {
115                assert!(api_base_url.is_none());
116                assert_eq!(label.as_deref(), Some("alice's rig"));
117                assert!(reset);
118            }
119            other => panic!("expected register, got {other:?}"),
120        }
121    }
122
123    #[test]
124    fn parses_bare_register() {
125        let cli = Cli::parse_from(["studio-worker", "register"]);
126        assert!(matches!(
127            cli.command,
128            Command::Register {
129                api_base_url: None,
130                label: None,
131                reset: false,
132            }
133        ));
134    }
135
136    #[test]
137    fn parses_set_threshold_with_float() {
138        let cli = Cli::parse_from(["studio-worker", "set-threshold", "12.5"]);
139        match cli.command {
140            Command::SetThreshold { gb } => assert!((gb - 12.5).abs() < 1e-6),
141            other => panic!("expected set-threshold, got {other:?}"),
142        }
143    }
144
145    #[test]
146    fn parses_all_simple_subcommands() {
147        let cases = [
148            ("status", Command::Status),
149            ("install-service", Command::InstallService),
150            ("uninstall-service", Command::UninstallService),
151            ("enable", Command::Enable),
152            ("disable", Command::Disable),
153            ("config", Command::Config),
154            ("check-update", Command::CheckUpdate),
155            ("ui", Command::Ui),
156        ];
157        for (name, expected) in cases {
158            let cli = Cli::parse_from(["studio-worker", name]);
159            assert_eq!(cli.command, expected, "parsing `{name}`");
160        }
161    }
162}