chore: support allocating testnets to the local DC

bazel/defs.bzl

@@ -368,6 +368,26 @@ write_info_file_var = rule(
     },
 )
 
+def _write_version_file_var_impl(ctx):
+    """Helper rule that creates a file with the content of the provided var from the version file."""
+
+    output = ctx.actions.declare_file(ctx.label.name)
+    ctx.actions.run_shell(
+        command = """
+            grep <{version_file} -e '{varname}' \\
+                    | cut -d' ' -f2 > {out}""".format(varname = ctx.attr.varname, version_file = ctx.version_file.path, out = output.path),
+        inputs = [ctx.version_file],
+        outputs = [output],
+    )
+    return [DefaultInfo(files = depset([output]))]
+
+write_version_file_var = rule(
+    implementation = _write_version_file_var_impl,
+    attrs = {
+        "varname": attr.string(mandatory = True),
+    },
+)
+
 def file_size_check(
         name,
         file,

bazel/workspace_status.sh

@@ -27,3 +27,7 @@ if [ -n "${CI_JOB_NAME:-}" ]; then
     STABLE_FARM_METADATA="$STABLE_FARM_METADATA;JOB_NAME=$CI_JOB_NAME"
 fi
 echo "STABLE_FARM_METADATA $STABLE_FARM_METADATA"
+
+# Used for allocating a Farm testnet to the local DC in CI (Search for allocate_testnet_to_local_dc)
+NODE_NAME="${NODE_NAME:-}"
+echo "DC ${NODE_NAME%%-*}"

rs/tests/BUILD.bazel

@@ -3,7 +3,7 @@ load("@rules_distroless//apt:defs.bzl", "dpkg_status")
 load("@rules_distroless//distroless:defs.bzl", "passwd")
 load("@rules_oci//oci:defs.bzl", "oci_image")
 load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
-load("//bazel:defs.bzl", "write_info_file_var")
+load("//bazel:defs.bzl", "write_info_file_var", "write_version_file_var")
 load(":system_tests.bzl", "oci_tar", "uvm_config_image")
 
 package(default_visibility = ["//rs:system-tests-pkg"])
@@ -284,3 +284,8 @@ write_info_file_var(
     varname = "STABLE_FARM_METADATA",
     visibility = ["//visibility:public"],
 )
+
+write_version_file_var(
+    name = "DC.txt",
+    varname = "DC",
+)

rs/tests/driver/src/driver/group.rs

@@ -648,6 +648,7 @@ impl SystemTestSubGroup {
 }
 
 pub struct SystemTestGroup {
+    allocate_testnet_to_local_dc: bool,
     setup: Option<Box<dyn PotSetupFn>>,
     teardowns: Vec<Box<dyn PotSetupFn>>,
     tests: Vec<SystemTestSubGroup>,
@@ -692,6 +693,7 @@ impl TestEnvAttribute for CliArguments {
 impl SystemTestGroup {
     pub fn new() -> Self {
         Self {
+            allocate_testnet_to_local_dc: false,
             setup: Default::default(),
             teardowns: Default::default(),
             tests: Default::default(),
@@ -744,6 +746,11 @@ impl SystemTestGroup {
         self
     }
 
+    pub fn allocate_testnet_to_local_dc(mut self) -> Self {
+        self.allocate_testnet_to_local_dc = true;
+        self
+    }
+
     pub fn with_setup<F: PotSetupFn>(mut self, setup: F) -> Self {
         self.setup = Some(Box::new(setup));
         self
@@ -1296,7 +1303,11 @@ impl SystemTestGroup {
             }
             InfraProvider::Farm.write_attribute(&root_env);
             if with_farm {
-                root_env.create_group_setup(group_ctx.group_base_name.clone(), args.no_group_ttl);
+                root_env.create_group_setup(
+                    group_ctx.group_base_name.clone(),
+                    self.allocate_testnet_to_local_dc,
+                    args.no_group_ttl,
+                );
             }
             debug!(group_ctx.log(), "Created group context: {:?}", group_ctx);
         }

rs/tests/driver/src/driver/test_env_api.rs

@@ -134,7 +134,7 @@
 use super::{
     config::NODES_INFO,
     driver_setup::SSH_AUTHORIZED_PRIV_KEYS_DIR,
-    farm::{DnsRecord, PlaynetCertificate},
+    farm::{DnsRecord, HostFeature, PlaynetCertificate},
     test_setup::{GroupSetup, InfraProvider},
 };
 use crate::{
@@ -1482,11 +1482,21 @@ pub fn get_build_setupos_config_image_tool() -> PathBuf {
 }
 
 pub trait HasGroupSetup {
-    fn create_group_setup(&self, group_base_name: String, no_group_ttl: bool);
+    fn create_group_setup(
+        &self,
+        group_base_name: String,
+        allocate_testnet_to_local_dc: bool,
+        no_group_ttl: bool,
+    );
 }
 
 impl HasGroupSetup for TestEnv {
-    fn create_group_setup(&self, group_base_name: String, no_group_ttl: bool) {
+    fn create_group_setup(
+        &self,
+        group_base_name: String,
+        allocate_testnet_to_local_dc: bool,
+        no_group_ttl: bool,
+    ) {
         let log = self.logger();
         if GroupSetup::attribute_exists(self) {
             let group_setup = GroupSetup::read_attribute(self);
@@ -1501,11 +1511,24 @@ impl HasGroupSetup for TestEnv {
             let group_setup = GroupSetup::new(group_base_name.clone(), timeout);
             match InfraProvider::read_attribute(self) {
                 InfraProvider::Farm => {
+                    let required_host_features = match allocate_testnet_to_local_dc
+                        .then(|| std::env::var("DC").expect("Expected env var 'DC' to be set"))
+                    {
+                        Some(dc) if !dc.is_empty() => vec![HostFeature::DC(dc)],
+                        _ => vec![],
+                    };
+                    info!(
+                        log,
+                        "Creating group {} with required_host_features: {:?} ...",
+                        group_setup.infra_group_name,
+                        required_host_features
+                    );
+
                     let farm_base_url = FarmBaseUrl::read_attribute(self);
                     let farm = Farm::new(farm_base_url.into(), self.logger());
                     let group_spec = GroupSpec {
                         vm_allocation: None,
-                        required_host_features: vec![],
+                        required_host_features,
                         preferred_network: None,
                         metadata: None,
                     };

rs/tests/nested/guestos_upgrade.rs

@@ -22,6 +22,7 @@ use nested::{
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(nested::setup)
         .add_test(systest!(upgrade_guestos))
         .with_timeout_per_test(Duration::from_secs(30 * 60))

rs/tests/nested/hostos_upgrade.rs

@@ -17,6 +17,7 @@ use nested::util::{
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(nested::setup)
         .add_test(systest!(upgrade_hostos))
         .with_timeout_per_test(Duration::from_secs(30 * 60))

rs/tests/nested/nns_recovery/nr_all_broken_seq_np_actions.rs

@@ -32,6 +32,7 @@ use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(|env| {
             setup(
                 env,

rs/tests/nested/nns_recovery/nr_broken_dfinity_node.rs

@@ -31,6 +31,7 @@ use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(|env| {
             setup(
                 env,

rs/tests/nested/nns_recovery/nr_local.rs

@@ -31,6 +31,7 @@ use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(|env| {
             setup(
                 env,

rs/tests/nested/nns_recovery/nr_no_bless_fix_like_np.rs

@@ -32,6 +32,7 @@ use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(|env| {
             setup(
                 env,

rs/tests/nested/registration.rs

@@ -4,6 +4,7 @@ use std::time::Duration;
 
 fn main() -> Result<()> {
     SystemTestGroup::new()
+        .allocate_testnet_to_local_dc()
         .with_setup(nested::setup)
         .add_test(systest!(nested::registration))
         .with_timeout_per_test(Duration::from_secs(20 * 60))

rs/tests/system_tests.bzl

@@ -133,6 +133,7 @@ def system_test(
     icos_images |= icos_config.icos_images
 
     env_var_files["FARM_METADATA"] = "//rs/tests:farm_metadata.txt"
+    env_var_files["DC"] = "//rs/tests:DC.txt"
 
     extra_args_simple = []
     extra_args_colocated = []