When bucketing decreasing in a rollout, then end enrollment

components/nimbus/src/enrollment.rs

@@ -237,6 +237,14 @@ impl ExperimentEnrollment {
                             out_enrollment_events.push(updated_enrollment.get_change_event());
                             updated_enrollment
                         }
+                        EnrollmentStatus::NotEnrolled {
+                            reason: NotEnrolledReason::NotSelected,
+                        } => {
+                            // In the case of a rollout being scaled back, we should end with WasEnrolled.
+                            //
+                            self.on_experiment_ended(out_enrollment_events)
+                                .ok_or_else(|| NimbusError::InternalError("An unexpected None happened while ending an experiment prematurely"))?
+                        }
                         EnrollmentStatus::NotEnrolled { .. }
                         | EnrollmentStatus::Enrolled { .. }
                         | EnrollmentStatus::Disqualified { .. }

components/nimbus/src/tests/helpers.rs

@@ -195,3 +195,46 @@ pub fn get_ios_rollout_experiment() -> Experiment {
     }))
     .unwrap()
 }
+
+pub fn get_bucketed_rollout(slug: &str, count: i64) -> Experiment {
+    let feature_id = "a-feature";
+    serde_json::from_value(json!(
+        {
+        "schemaVersion": "1.0.0",
+        "slug": slug,
+        "endDate": null,
+        "branches":[
+            {
+                "slug": "control",
+                "ratio": 1,
+                "feature": {
+                    "featureId": feature_id,
+                    "enabled": true,
+                    "value": {},
+                }
+            },
+        ],
+        "isRollout": true,
+        "featureIds": [feature_id],
+        "channel": "nightly",
+        "probeSets": [],
+        "startDate": null,
+        "appName":"fenix",
+        "appId":"org.mozilla.fenix",
+        "bucketConfig":{
+            // Also enroll everyone.
+            "count": count,
+            "start": 0,
+            "total": 10_000,
+            "namespace":"secure-silver",
+            "randomizationUnit":"nimbus_id"
+        },
+        "userFacingName":"",
+        "referenceBranch":"control",
+        "isEnrollmentPaused":false,
+        "proposedEnrollment":7,
+        "userFacingDescription":"",
+    }
+    ))
+    .unwrap()
+}

components/nimbus/src/tests/test_enrollment.rs

@@ -4,6 +4,7 @@
 
 // Testing enrollment.rs
 
+use crate::tests::helpers::get_bucketed_rollout;
 use crate::{
     defaults::Defaults,
     enrollment::*,
@@ -775,8 +776,7 @@ fn test_evolver_experiment_update_enrolled_then_targeting_changed() -> Result<()
 
 #[test]
 fn test_evolver_experiment_update_enrolled_then_bucketing_changed() -> Result<()> {
-    let mut exp = get_test_experiments()[0].clone();
-    exp.bucket_config.count = 0; // Make the experiment bucketing fail.
+    let exp = get_bucketed_rollout("test-rollout", 0);
     let (nimbus_id, app_ctx, aru) = local_ctx();
     let th = app_ctx.into();
     let evolver = enrollment_evolver(&nimbus_id, &th, &aru);
@@ -790,7 +790,7 @@ fn test_evolver_experiment_update_enrolled_then_bucketing_changed() -> Result<()
             reason: EnrolledReason::Qualified,
         },
     };
-    let enrollment = evolver
+    let observed = evolver
         .evolve_enrollment(
             true,
             Some(&exp),
@@ -799,8 +799,68 @@ fn test_evolver_experiment_update_enrolled_then_bucketing_changed() -> Result<()
             &mut events,
         )?
         .unwrap();
-    assert_eq!(enrollment, existing_enrollment);
-    assert!(events.is_empty());
+    assert!(matches!(
+        observed,
+        ExperimentEnrollment {
+            status: EnrollmentStatus::WasEnrolled { .. },
+            ..
+        }
+    ));
+    assert_eq!(1, events.len());
+    Ok(())
+}
+
+#[test]
+fn test_rollout_unenrolls_when_bucketing_changes() -> Result<()> {
+    let (nimbus_id, app_ctx, aru) = local_ctx();
+    let th = app_ctx.into();
+    let evolver = enrollment_evolver(&nimbus_id, &th, &aru);
+
+    let slug = "my-rollout";
+
+    // Start at 0%
+    let ro = get_bucketed_rollout(slug, 0);
+    let recipes = [ro];
+
+    let (enrollments, _) = evolver.evolve_enrollments::<Experiment>(true, &[], &recipes, &[])?;
+
+    assert_eq!(enrollments.len(), 1);
+    let enr = enrollments.get(0).unwrap();
+    assert_eq!(&enr.slug, slug);
+    assert!(matches!(&enr.status, EnrollmentStatus::NotEnrolled { .. }));
+
+    // Up to 100%
+    let prev_recipes = recipes;
+    let ro = get_bucketed_rollout(slug, 10_000);
+    let recipes = [ro];
+
+    let (enrollments, _) = evolver.evolve_enrollments::<Experiment>(
+        true,
+        &prev_recipes,
+        &recipes,
+        enrollments.as_slice(),
+    )?;
+    assert_eq!(enrollments.len(), 1);
+    let enr = enrollments.get(0).unwrap();
+    assert_eq!(&enr.slug, slug);
+    assert!(matches!(&enr.status, EnrollmentStatus::Enrolled { .. }));
+
+    // Back to zero again
+    let prev_recipes = recipes;
+    let ro = get_bucketed_rollout(slug, 0);
+    let recipes = [ro];
+
+    let (enrollments, _) = evolver.evolve_enrollments::<Experiment>(
+        true,
+        &prev_recipes,
+        &recipes,
+        enrollments.as_slice(),
+    )?;
+    assert_eq!(enrollments.len(), 1);
+    let enr = enrollments.get(0).unwrap();
+    assert_eq!(&enr.slug, slug);
+    assert!(matches!(&enr.status, EnrollmentStatus::WasEnrolled { .. }));
+
     Ok(())
 }