Merge pull request #253 from ptournaris-ai/feat/site-flag-on-auth-status

feat(auth): add --site flag to auth status command

Author: platinummonkey

src/main.rs

@@ -5610,7 +5610,12 @@ enum AuthActions {
     /// Logout and clear tokens
     Logout,
     /// Check authentication status
-    Status,
+    Status {
+        /// Datadog site to check status for (e.g. datadoghq.eu, us3.datadoghq.com).
+        /// Overrides DD_SITE env var and config file. Defaults to datadoghq.com.
+        #[arg(long, value_name = "SITE")]
+        site: Option<String>,
+    },
     /// Print access token (debug builds only)
     #[cfg(debug_assertions)]
     Token,
@@ -8194,7 +8199,12 @@ async fn main_inner() -> anyhow::Result<()> {
                 commands::auth::login(&cfg, resolved).await?
             }
             AuthActions::Logout => commands::auth::logout(&cfg).await?,
-            AuthActions::Status => commands::auth::status(&cfg)?,
+            AuthActions::Status { site } => {
+                if let Some(s) = site {
+                    cfg.site = s;
+                }
+                commands::auth::status(&cfg)?
+            }
             #[cfg(debug_assertions)]
             AuthActions::Token => commands::auth::token(&cfg)?,
             AuthActions::Refresh => commands::auth::refresh(&cfg).await?,

src/test_commands.rs

@@ -3170,3 +3170,43 @@ async fn test_llm_obs_spans_search_no_auth() {
     assert!(result.is_err(), "should fail without auth");
     cleanup_env();
 }
+
+// -------------------------------------------------------------------------
+// Auth status --site flag
+// -------------------------------------------------------------------------
+
+#[test]
+fn test_auth_status_accepts_site_flag() {
+    use clap::Parser;
+
+    let cli = crate::Cli::try_parse_from(["pup", "auth", "status", "--site", "datadoghq.eu"])
+        .expect("auth status --site should parse");
+
+    match cli.command {
+        crate::Commands::Auth { action } => match action {
+            crate::AuthActions::Status { site } => {
+                assert_eq!(site, Some("datadoghq.eu".to_string()));
+            }
+            _ => panic!("expected AuthActions::Status"),
+        },
+        _ => panic!("expected Commands::Auth"),
+    }
+}
+
+#[test]
+fn test_auth_status_site_flag_is_optional() {
+    use clap::Parser;
+
+    let cli = crate::Cli::try_parse_from(["pup", "auth", "status"])
+        .expect("auth status without --site should parse");
+
+    match cli.command {
+        crate::Commands::Auth { action } => match action {
+            crate::AuthActions::Status { site } => {
+                assert_eq!(site, None);
+            }
+            _ => panic!("expected AuthActions::Status"),
+        },
+        _ => panic!("expected Commands::Auth"),
+    }
+}