(Chore): Better telemetry for quarantine query

bundle/src/types.rs

@@ -28,7 +28,7 @@
 }
 
 impl Test {
     // trunk-ignore(clippy/too_many_arguments)
     pub fn new<T: AsRef<str> + ToString>(
         id: Option<T>,
         name: String,
@@ -104,6 +104,8 @@
     pub group_is_quarantined: bool,
     #[serde(rename = "quarantineResults")]
     pub quarantine_results: Vec<Test>,
+    #[serde(rename = "quarantiningDisabledForRepo")]
+    pub quarantining_disabled_for_repo: bool,
 }
 
 #[derive(Debug, Serialize, Clone)]

cli/src/context_quarantine.rs

@@ -21,6 +21,7 @@ use context::{
 };
 use pluralizer::pluralize;
 use prost::Message;
+use proto::upload_metrics::trunk::QuarantineQueryResult;
 
 #[derive(Debug)]
 pub enum QuarantineFetchStatus {
@@ -71,6 +72,21 @@ impl QuarantineContext {
     }
 }
 
+pub fn quarantine_query_result(
+    disable_quarantining: bool,
+    ctx: &QuarantineContext,
+) -> proto::upload_metrics::trunk::QuarantineQueryResult {
+    if disable_quarantining || ctx.quarantine_status.quarantining_disabled_for_repo {
+        return QuarantineQueryResult::Disabled;
+    }
+
+    match &ctx.fetch_status {
+        QuarantineFetchStatus::FetchFailed(_) => QuarantineQueryResult::Failure,
+        QuarantineFetchStatus::FetchSkipped => QuarantineQueryResult::Skipped,
+        QuarantineFetchStatus::FetchSucceeded => QuarantineQueryResult::Success,
+    }
+}
+
 fn convert_case_to_test<T: AsRef<str> + ToString>(
     repo: &RepoUrlParts,
     org_slug: T,
@@ -334,7 +350,10 @@ pub async fn gather_quarantine_context(
         tracing::info!("Quarantining is not enabled, not quarantining any tests");
         return Ok(QuarantineContext {
             exit_code,
-            quarantine_status: QuarantineBulkTestStatus::default(),
+            quarantine_status: QuarantineBulkTestStatus {
+                quarantining_disabled_for_repo: true,
+                ..Default::default()
+            },
             failures: failed_tests_extractor.failed_tests().to_vec(),
             repo: request.repo.clone(),
             org_url_slug: request.org_url_slug.clone(),
@@ -627,6 +646,53 @@ mod tests {
         assert_eq!(multi_failures[0].name, "Hello");
     }
 
+    #[test]
+    fn test_quarantine_query_result_mapper() {
+        use proto::upload_metrics::trunk::QuarantineQueryResult;
+
+        let success_ctx = QuarantineContext {
+            exit_code: 0,
+            quarantine_status: QuarantineBulkTestStatus::default(),
+            failures: vec![],
+            repo: RepoUrlParts::default(),
+            org_url_slug: String::new(),
+            fetch_status: QuarantineFetchStatus::FetchSucceeded,
+        };
+        assert_eq!(
+            quarantine_query_result(false, &success_ctx),
+            QuarantineQueryResult::Success
+        );
+        assert_eq!(
+            quarantine_query_result(true, &success_ctx),
+            QuarantineQueryResult::Disabled
+        );
+
+        let skipped_ctx = QuarantineContext {
+            fetch_status: QuarantineFetchStatus::FetchSkipped,
+            ..success_ctx
+        };
+        assert_eq!(
+            quarantine_query_result(false, &skipped_ctx),
+            QuarantineQueryResult::Skipped
+        );
+
+        let repo_disabled_ctx = QuarantineContext {
+            quarantine_status: QuarantineBulkTestStatus {
+                quarantining_disabled_for_repo: true,
+                ..Default::default()
+            },
+            fetch_status: QuarantineFetchStatus::FetchSucceeded,
+            exit_code: 0,
+            failures: vec![],
+            repo: RepoUrlParts::default(),
+            org_url_slug: String::new(),
+        };
+        assert_eq!(
+            quarantine_query_result(false, &repo_disabled_ctx),
+            QuarantineQueryResult::Disabled
+        );
+    }
+
     #[test]
     fn test_convert_case_to_test_populates_failure_message() {
         use quick_junit::{NonSuccessKind, TestCase, TestCaseStatus, TestSuite};

cli/src/main.rs

@@ -192,7 +192,7 @@
     );
     match cli.command {
         Commands::Upload(upload_args) => {
-            let result = run_upload(upload_args, None, None, Some(render_sender)).await?;
+            let result = run_upload(upload_args, None, None, Some(render_sender), None).await?;
             Ok(RunResult::Upload(result))
         }
         Commands::Test(test_args) => {
@@ -225,7 +225,7 @@
 ) -> anyhow::Result<()> {
     let command_string = String::from(command_name);
     let sentry_layer = sentry_tracing::layer().event_mapper(move |event, context| {
         // trunk-ignore(clippy/match_ref_pats)
         match event.metadata().level() {
             &tracing::Level::ERROR => {
                 let mut event = sentry_tracing::event_from_event(event, context);

cli/src/test_command.rs

@@ -1,24 +1,24 @@
 use std::{
     env,
     process::{Command, Stdio},
-    sync::{mpsc::Sender, Arc},
+    sync::{Arc, mpsc::Sender},
     time::SystemTime,
 };
 
 use clap::Args;
 use constants::EXIT_FAILURE;
 use display::{
     end_output::EndOutput,
-    message::{send_message, DisplayMessage},
+    message::{DisplayMessage, send_message},
 };
 use superconsole::{
-    style::{Attribute, Stylize},
     Line, Span,
+    style::{Attribute, Stylize},
 };
 
 use crate::{
     context::{gather_debug_props, gather_initial_test_context},
-    upload_command::{run_upload, UploadArgs, UploadRunResult},
+    upload_command::{UploadArgs, UploadRunResult, run_upload},
 };
 
 enum RunOutput {
@@ -136,6 +136,7 @@ pub async fn run_test(
         Some(test_context),
         Some(test_run_result.clone()),
         None,
+        None,
     )
     .await;
     match upload_run_result {

cli/src/upload_command.rs

@@ -18,7 +18,7 @@ use superconsole::{
 };
 use tempfile::TempDir;
 
-use crate::context_quarantine::QuarantineContext;
+use crate::context_quarantine::{QuarantineContext, quarantine_query_result};
 use crate::validate_command::JunitReportValidations;
 use crate::{
     context::{
@@ -360,6 +360,7 @@ pub async fn run_upload(
     pre_test_context: Option<PreTestContext>,
     test_run_result: Option<TestRunResult>,
     render_sender: Option<Sender<DisplayMessage>>,
+    quarantine_query_result_override: Option<proto::upload_metrics::trunk::QuarantineQueryResult>,
 ) -> anyhow::Result<UploadRunResult> {
     // grab the exec start if provided (`test` subcommand) or use the current time
     let cli_started_at = if let Some(test_run_result) = test_run_result.as_ref() {
@@ -509,6 +510,8 @@ pub async fn run_upload(
         upload_args.dry_run,
     )
     .await;
+    let quarantine_query_result = quarantine_query_result_override
+        .unwrap_or_else(|| quarantine_query_result(disable_quarantining, &quarantine_context));
     let upload_metrics = proto::upload_metrics::trunk::UploadMetrics {
         client_version: Some(proto::upload_metrics::trunk::Semver {
             major: env!("CARGO_PKG_VERSION_MAJOR").parse().unwrap_or_default(),
@@ -526,6 +529,7 @@ pub async fn run_upload(
         upload_finished_at: Some(chrono::Utc::now().into()),
         failed: false,
         failure_reason: "".into(),
+        quarantine_query_result: quarantine_query_result.into(),
     };
     let mut request = api::message::TelemetryUploadMetricsRequest { upload_metrics };
     if !upload_args.dry_run {

cli/tests/upload.rs

@@ -34,6 +34,7 @@ use lazy_static::lazy_static;
 use predicates::prelude::*;
 use pretty_assertions::assert_eq;
 use prost::Message;
+use proto::upload_metrics::trunk::QuarantineQueryResult;
 use tempfile::tempdir;
 #[cfg(target_os = "macos")]
 use test_utils::inputs::unpack_archive_to_dir;
@@ -1277,6 +1278,10 @@ async fn telemetry_upload_metrics_on_upload_failure() {
     assert_eq!(telemetry_request_repo.owner, "trunk-io");
     assert_eq!(telemetry_request_repo.name, "analytics-cli");
     assert_eq!(telemetry_request.failure_reason, "400_bad_request");
+    assert_eq!(
+        QuarantineQueryResult::try_from(telemetry_request.quarantine_query_result).unwrap(),
+        QuarantineQueryResult::Disabled
+    );
 
     // HINT: View CLI output with `cargo test -- --nocapture`
     println!("{assert}");
@@ -1308,11 +1313,126 @@ async fn telemetry_upload_metrics_on_upload_success() {
     assert_eq!(telemetry_request_repo.owner, "trunk-io");
     assert_eq!(telemetry_request_repo.name, "analytics-cli");
     assert_eq!(telemetry_request.failure_reason, "");
+    assert_eq!(
+        QuarantineQueryResult::try_from(telemetry_request.quarantine_query_result).unwrap(),
+        QuarantineQueryResult::Disabled
+    );
 
     // HINT: View CLI output with `cargo test -- --nocapture`
     println!("{assert}");
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn telemetry_upload_metrics_quarantine_query_result_success() {
+    let temp_dir = tempdir().unwrap();
+    generate_mock_git_repo(&temp_dir);
+    generate_mock_valid_junit_xmls_with_failures(&temp_dir);
+
+    let state = MockServerBuilder::new().spawn_mock_server().await;
+
+    CommandBuilder::upload(temp_dir.path(), state.host.clone())
+        .command()
+        .assert()
+        .failure();
+
+    let requests = state.requests.lock().unwrap().clone();
+    let quarantine_query_result = requests
+        .iter()
+        .rev()
+        .find_map(|req| match req {
+            RequestPayload::TelemetryUploadMetrics(ur) => Some(ur.quarantine_query_result),
+            _ => None,
+        })
+        .expect("telemetry upload metrics request");
+    assert_eq!(
+        QuarantineQueryResult::try_from(quarantine_query_result).unwrap(),
+        QuarantineQueryResult::Success
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn telemetry_upload_metrics_quarantine_query_result_failure() {
+    let temp_dir = tempdir().unwrap();
+    generate_mock_git_repo(&temp_dir);
+    generate_mock_valid_junit_xmls_with_failures(&temp_dir);
+
+    let mut mock_server_builder = MockServerBuilder::new();
+    mock_server_builder.set_get_quarantining_config_handler(
+        |Json(_): Json<GetQuarantineConfigRequest>| async {
+            Err::<Json<GetQuarantineConfigResponse>, StatusCode>(StatusCode::INTERNAL_SERVER_ERROR)
+        },
+    );
+    let state = mock_server_builder.spawn_mock_server().await;
+
+    CommandBuilder::upload(temp_dir.path(), state.host.clone())
+        .command()
+        .assert()
+        .failure();
+
+    let requests = state.requests.lock().unwrap().clone();
+    let quarantine_query_result = requests
+        .iter()
+        .rev()
+        .find_map(|req| match req {
+            RequestPayload::TelemetryUploadMetrics(ur) => Some(ur.quarantine_query_result),
+            _ => None,
+        })
+        .expect("telemetry upload metrics request");
+    assert_eq!(
+        QuarantineQueryResult::try_from(quarantine_query_result).unwrap(),
+        QuarantineQueryResult::Failure
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn telemetry_upload_metrics_quarantine_query_result_skipped() {
+    let temp_dir = tempdir().unwrap();
+    generate_mock_git_repo(&temp_dir);
+
+    // All-passing JUnit so quarantine lookup is skipped (no failed tests).
+    let test_case_options = junit_mock::TestCaseOptions {
+        test_case_names: Some(vec![String::from("test_case")]),
+        test_case_classnames: Some(vec![String::from("TestClass")]),
+        test_case_random_count: 0usize,
+        test_case_sys_out_percentage: 0u8,
+        test_case_sys_err_percentage: 0u8,
+        test_case_duration_range: vec![Duration::new(10, 0).into(), Duration::new(20, 0).into()],
+        test_case_success_to_skip_to_fail_to_error_percentage: vec![vec![100u8, 0u8, 0u8, 0u8]],
+    };
+    let options = junit_mock::Options {
+        global: junit_mock::GlobalOptions::try_parse_from([""]).unwrap(),
+        report: junit_mock::ReportOptions::try_parse_from([""]).unwrap(),
+        test_suite: junit_mock::TestSuiteOptions::try_parse_from([""]).unwrap(),
+        test_case: test_case_options,
+        test_rerun: junit_mock::TestRerunOptions::try_parse_from([""]).unwrap(),
+    };
+    let mut mock = JunitMock::new(options);
+    let reports = mock.generate_reports();
+    mock.write_reports_to_file(temp_dir.as_ref(), reports)
+        .unwrap();
+
+    let state = MockServerBuilder::new().spawn_mock_server().await;
+
+    CommandBuilder::upload(temp_dir.path(), state.host.clone())
+        .command()
+        .assert()
+        .success();
+
+    let requests = state.requests.lock().unwrap().clone();
+    let quarantine_query_result = requests
+        .iter()
+        .rev()
+        .find_map(|req| match req {
+            RequestPayload::TelemetryUploadMetrics(ur) => Some(ur.quarantine_query_result),
+            _ => None,
+        })
+        .expect("telemetry upload metrics request");
+    assert_eq!(
+        QuarantineQueryResult::try_from(quarantine_query_result).unwrap(),
+        QuarantineQueryResult::Skipped
+    );
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn telemetry_does_not_impact_return() {
     let temp_dir = tempdir().unwrap();

context/src/junit/bindings/test_case.rs

@@ -332,7 +332,7 @@ impl TryInto<TestCase> for BindingsTestCase {
             properties,
             bazel_run_information: _,
         } = self;
-        // donotland: anything here?
+
         let mut test_case = TestCase::new(name, status.try_into()?);
         test_case.classname = classname.map(|c| c.into());
         test_case.assertions = assertions;

proto/proto/upload_metrics.proto

@@ -5,6 +5,16 @@ package trunk.oss.flakytests_cli.v1;
 import "common.proto";
 import "google/protobuf/timestamp.proto";
 
+enum QuarantineQueryResult {
+  QUARANTINE_QUERY_RESULT_UNSPECIFIED = 0;
+  QUARANTINE_QUERY_RESULT_SUCCESS = 1;
+  QUARANTINE_QUERY_RESULT_FAILURE = 2;
+  QUARANTINE_QUERY_RESULT_DISABLED = 3;
+  QUARANTINE_QUERY_RESULT_SKIPPED = 4;
+  // rspec only
+  QUARANTINE_QUERY_RESULT_CACHED = 5;
+}
+
 message UploadMetrics {
   Semver client_version = 1;
   Repo repo = 2;
@@ -13,6 +23,7 @@ message UploadMetrics {
   google.protobuf.Timestamp upload_finished_at = 5;
   bool failed = 6;
   string failure_reason = 7;
+  QuarantineQueryResult quarantine_query_result = 8;
 }
 
 // Used by the analytics-uploader, kept here to avoid collisions in the future.