chore: support allocating testnets to the local DC

bazel/BUILD.bazel

@@ -1,6 +1,6 @@
 load("@bazel_skylib//rules:common_settings.bzl", "bool_flag", "string_flag")
 load("@buildifier_prebuilt//:rules.bzl", "buildifier")
-load("//bazel:defs.bzl", "write_info_file_var")
+load("//bazel:defs.bzl", "volatile_status", "write_info_file_var")
 
 bool_flag(
     name = "enable_malicious_code",
@@ -72,6 +72,11 @@ write_info_file_var(
     visibility = ["//visibility:public"],
 )
 
+volatile_status(
+    name = "volatile-status.txt",
+    visibility = ["//visibility:public"],
+)
+
 exports_files(
     [
         "prost_generator.sh",

bazel/defs.bzl

@@ -368,6 +368,16 @@ write_info_file_var = rule(
     },
 )
 
+def _volatile_status_impl(ctx):
+    return [DefaultInfo(
+        files = depset([ctx.version_file]),
+        runfiles = ctx.runfiles(files = [ctx.version_file]),
+    )]
+
+volatile_status = rule(
+    implementation = _volatile_status_impl,
+)
+
 def file_size_check(
         name,
         file,

bazel/workspace_status.sh

@@ -27,3 +27,6 @@ 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)
+echo "NODE_NAME ${NODE_NAME:-}"

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,33 @@ 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 = if allocate_testnet_to_local_dc {
+                        let node_name = read_var_from_volatile_status_file("NODE_NAME");
+                        if let Some(name) = node_name {
+                            let dc = name
+                                .split_once('-')
+                                .expect("NODE_NAME doesn't contain '-'")
+                                .0
+                                .to_string();
+                            vec![HostFeature::DC(dc.clone())]
+                        } else {
+                            vec![]
+                        }
+                    } else {
+                        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,
                     };
@@ -1612,6 +1644,21 @@ pub fn read_dependency_from_env_to_string(v: &str) -> Result<String> {
     read_dependency_to_string(path_from_env)
 }
 
+pub fn read_var_from_volatile_status_file(var_name: &str) -> Option<String> {
+    let volatile_status_path = get_dependency_path_from_env("VOLATILE_STATUS_FILE");
+    let content = fs::read_to_string(&volatile_status_path).unwrap_or_else(|e| {
+        panic!("Couldn't read content of the {volatile_status_path:?} file: {e:?}")
+    });
+    for line in content.lines() {
+        if let Some((name, value)) = line.split_once(' ')
+            && name.trim() == var_name
+        {
+            return Some(value.trim().to_string());
+        }
+    }
+    None
+}
+
 pub fn load_wasm<P: AsRef<Path>>(p: P) -> Vec<u8> {
     let mut wasm_bytes = std::fs::read(get_dependency_path(&p))
         .unwrap_or_else(|e| panic!("Could not read WASM from {:?}: {e:?}", p.as_ref()));

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

@@ -121,6 +121,7 @@ def system_test(
     _runtime_deps = dict(runtime_deps)
 
     _runtime_deps["TEST_BIN"] = test_driver_target
+    _runtime_deps["VOLATILE_STATUS_FILE"] = "//bazel:volatile-status.txt"
 
     env_var_files = {}
     icos_images = dict()