test_default_cockroach_admin_addrs_from_blueprint()

Author: jgallagher

nexus/src/app/background/tasks/crdb_node_id_collector.rs

@@ -241,17 +241,15 @@ mod tests {
     use httptest::responders::status_code;
     use nexus_db_queries::db::pub_test_utils::TestDatabase;
     use nexus_reconfigurator_planning::blueprint_builder::BlueprintBuilder;
-    use nexus_sled_agent_shared::inventory::OmicronZoneDataset;
-    use nexus_types::deployment::BlueprintZoneConfig;
+    use nexus_reconfigurator_planning::example::ExampleSystemBuilder;
+    use nexus_reconfigurator_planning::planner::PlannerRng;
+    use nexus_types::deployment::BlueprintSource;
     use nexus_types::deployment::BlueprintZoneDisposition;
     use nexus_types::deployment::BlueprintZoneImageSource;
     use omicron_common::api::external::Generation;
-    use omicron_common::zpool_name::ZpoolName;
     use omicron_test_utils::dev;
-    use omicron_uuid_kinds::SledUuid;
-    use omicron_uuid_kinds::ZpoolUuid;
     use std::collections::BTreeMap;
-    use std::iter;
+    use std::collections::BTreeSet;
     use std::net::SocketAddr;
 
     // The `CockroachAdminFromBlueprintViaFixedPort` type above is the standard
@@ -260,109 +258,83 @@ mod tests {
     // can _write_ that test), so test it in isolation here.
     #[test]
     fn test_default_cockroach_admin_addrs_from_blueprint() {
-        // Construct an empty blueprint with one sled.
-        let sled_id = SledUuid::new_v4();
-        let mut blueprint = BlueprintBuilder::build_empty_with_sleds(
-            iter::once(sled_id),
-            "test",
-        );
-        let bp_sled = &mut blueprint
-            .sleds
-            .get_mut(&sled_id)
-            .expect("found entry for test sled");
-
-        let zpool_id = ZpoolUuid::new_v4();
-        let make_crdb_zone_config =
-            |disposition, id, addr: SocketAddrV6| BlueprintZoneConfig {
-                disposition,
-                id,
-                filesystem_pool: ZpoolName::new_external(zpool_id),
-                zone_type: BlueprintZoneType::CockroachDb(
-                    blueprint_zone_type::CockroachDb {
-                        address: addr,
-                        dataset: OmicronZoneDataset {
-                            pool_name: format!("oxp_{}", zpool_id)
-                                .parse()
-                                .unwrap(),
-                        },
-                    },
-                ),
-                image_source: BlueprintZoneImageSource::InstallDataset,
-            };
-
-        // Add three CRDB zones with known addresses; the first and third are
-        // in service, and the second is expunged. Only the first and third
-        // should show up when we ask for addresses below.
-        let crdb_id1 = OmicronZoneUuid::new_v4();
-        let crdb_id2 = OmicronZoneUuid::new_v4();
-        let crdb_id3 = OmicronZoneUuid::new_v4();
-        let crdb_addr1: SocketAddrV6 = "[2001:db8::1]:1111".parse().unwrap();
-        let crdb_addr2: SocketAddrV6 = "[2001:db8::2]:1234".parse().unwrap();
-        let crdb_addr3: SocketAddrV6 = "[2001:db8::3]:1234".parse().unwrap();
-        bp_sled
-            .zones
-            .insert_unique(make_crdb_zone_config(
-                BlueprintZoneDisposition::InService,
-                crdb_id1,
-                crdb_addr1,
-            ))
-            .expect("freshly generated zone IDs are unique");
-        bp_sled
-            .zones
-            .insert_unique(make_crdb_zone_config(
-                BlueprintZoneDisposition::Expunged {
+        const TEST_NAME: &str =
+            "test_default_cockroach_admin_addrs_from_blueprint";
+        let logctx = dev::test_setup_log(TEST_NAME);
+        let log = &logctx.log;
+
+        // Build an example system with one sled.
+        let (example_system, bp0) =
+            ExampleSystemBuilder::new(log, TEST_NAME).nsleds(1).build();
+        let input = example_system.input;
+
+        // `ExampleSystemBuilder` doesn't place any cockroach nodes; assert so
+        // we bail out early if that changes.
+        let ncockroach = bp0
+            .all_omicron_zones(BlueprintZoneDisposition::is_in_service)
+            .filter(|(_, z)| z.zone_type.is_cockroach())
+            .count();
+        assert_eq!(ncockroach, 0);
+
+        // This blueprint has no cockroach zones, so should have no admin addrs
+        // either.
+        let admin_addrs = CockroachAdminFromBlueprintViaFixedPort
+            .cockroach_admin_addrs(&bp0)
+            .collect::<Vec<_>>();
+        assert_eq!(admin_addrs, Vec::new());
+
+        // Add 5 cockroach zones to our sled.
+        let mut builder = BlueprintBuilder::new_based_on(
+            log,
+            &bp0,
+            &input,
+            TEST_NAME,
+            PlannerRng::from_entropy(),
+        )
+        .expect("constructed builder");
+        let sled_id = bp0.sleds().next().expect("1 sled");
+        for _ in 0..5 {
+            builder
+                .sled_add_zone_cockroachdb(
+                    sled_id,
+                    BlueprintZoneImageSource::InstallDataset,
+                )
+                .expect("added cockroach");
+        }
+        let mut bp1 = builder.build(BlueprintSource::Test);
+
+        // Mutate this blueprint: expunge 2 of the 5 zones. Record the expected
+        // admin addrs from the other three.
+        let mut expected = BTreeSet::new();
+        let sled_config = bp1.sleds.get_mut(&sled_id).unwrap();
+        let mut seen = 0;
+        for mut zone in sled_config.zones.iter_mut() {
+            if !zone.zone_type.is_cockroach() {
+                continue;
+            }
+            // Keep even; expunge odd.
+            if seen % 2 == 0 {
+                let ip = zone.underlay_ip();
+                let addr = SocketAddrV6::new(ip, COCKROACH_ADMIN_PORT, 0, 0);
+                expected.insert((zone.id, addr));
+            } else {
+                zone.disposition = BlueprintZoneDisposition::Expunged {
                     as_of_generation: Generation::new(),
                     ready_for_cleanup: false,
-                },
-                crdb_id2,
-                crdb_addr2,
-            ))
-            .expect("freshly generated zone IDs are unique");
-        bp_sled
-            .zones
-            .insert_unique(make_crdb_zone_config(
-                BlueprintZoneDisposition::InService,
-                crdb_id3,
-                crdb_addr3,
-            ))
-            .expect("freshly generated zone IDs are unique");
-
-        // Also add a non-CRDB zone to ensure it's filtered out.
-        bp_sled
-            .zones
-            .insert_unique(BlueprintZoneConfig {
-                disposition: BlueprintZoneDisposition::InService,
-                id: OmicronZoneUuid::new_v4(),
-                filesystem_pool: ZpoolName::new_external(ZpoolUuid::new_v4()),
-                zone_type: BlueprintZoneType::CruciblePantry(
-                    blueprint_zone_type::CruciblePantry {
-                        address: "[::1]:0".parse().unwrap(),
-                    },
-                ),
-                image_source: BlueprintZoneImageSource::InstallDataset,
-            })
-            .expect("freshly generated zone IDs are unique");
-
-        // We expect to see CRDB zones 1 and 3 with their IPs but the ports
-        // changed to `COCKROACH_ADMIN_PORT`.
-        let mut expected = vec![
-            (
-                crdb_id1,
-                SocketAddrV6::new(*crdb_addr1.ip(), COCKROACH_ADMIN_PORT, 0, 0),
-            ),
-            (
-                crdb_id3,
-                SocketAddrV6::new(*crdb_addr3.ip(), COCKROACH_ADMIN_PORT, 0, 0),
-            ),
-        ];
-        // We sort starting with zone id, since the original zones are sorted
-        // that way in a map.
-        expected.sort_unstable();
+                };
+            }
+            seen += 1;
+        }
+        assert_eq!(seen, 5);
+        assert_eq!(expected.len(), 3);
 
+        // Confirm we only see the three expected addrs.
         let admin_addrs = CockroachAdminFromBlueprintViaFixedPort
-            .cockroach_admin_addrs(&blueprint)
-            .collect::<Vec<_>>();
+            .cockroach_admin_addrs(&bp1)
+            .collect::<BTreeSet<_>>();
         assert_eq!(expected, admin_addrs);
+
+        logctx.cleanup_successful();
     }
 
     #[tokio::test]