From 0121dec1e98bf0562243876ce7a11667618502f7 Mon Sep 17 00:00:00 2001
From: Eero Kelly <eero.kelly@dfinity.org>
Date: Fri, 10 Jul 2026 23:02:17 +0000
Subject: [PATCH] Normalize use of version in TopologySnapshot

---
 rs/tests/driver/src/driver/test_env_api.rs | 14 +++++++-------
 rs/tests/nns/node_swap_test.rs             |  2 +-
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/rs/tests/driver/src/driver/test_env_api.rs b/rs/tests/driver/src/driver/test_env_api.rs
index 69c8dcac4225..66b3e274d280 100644
--- a/rs/tests/driver/src/driver/test_env_api.rs
+++ b/rs/tests/driver/src/driver/test_env_api.rs
@@ -399,7 +399,7 @@ impl TopologySnapshot {
     }
 
     pub fn subnets(&self) -> Box<dyn Iterator<Item = SubnetSnapshot>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
         Box::new(
             self.local_registry
                 .get_subnet_ids(registry_version)
@@ -418,7 +418,7 @@ impl TopologySnapshot {
     }
 
     pub fn subnet_canister_ranges(&self, sub: SubnetId) -> Vec<CanisterIdRange> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
         self.local_registry
             .get_subnet_canister_ranges(registry_version, sub)
             .expect("Could not deserialize optional routing table from local registry.")
@@ -426,7 +426,7 @@ impl TopologySnapshot {
     }
 
     pub fn unassigned_nodes(&self) -> Box<dyn Iterator<Item = IcNodeSnapshot>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
         let assigned_nodes: HashSet<_> = self
             .local_registry
             .get_subnet_ids(registry_version)
@@ -468,7 +468,7 @@ impl TopologySnapshot {
     }
 
     pub fn api_boundary_nodes(&self) -> Box<dyn Iterator<Item = IcNodeSnapshot>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
 
         Box::new(
             self.local_registry
@@ -488,7 +488,7 @@ impl TopologySnapshot {
     }
 
     pub fn system_api_boundary_nodes(&self) -> Box<dyn Iterator<Item = IcNodeSnapshot>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
 
         Box::new(
             self.local_registry
@@ -508,7 +508,7 @@ impl TopologySnapshot {
     }
 
     pub fn app_api_boundary_nodes(&self) -> Box<dyn Iterator<Item = IcNodeSnapshot>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
 
         Box::new(
             self.local_registry
@@ -528,7 +528,7 @@ impl TopologySnapshot {
     }
 
     pub fn replica_version_records(&self) -> Result<Vec<(String, ReplicaVersionRecord)>> {
-        let registry_version = self.local_registry.get_latest_version();
+        let registry_version = self.registry_version;
 
         self.local_registry
             .get_all_replica_version_records(registry_version)?
diff --git a/rs/tests/nns/node_swap_test.rs b/rs/tests/nns/node_swap_test.rs
index 41ea8ea31ad0..a353d240185e 100644
--- a/rs/tests/nns/node_swap_test.rs
+++ b/rs/tests/nns/node_swap_test.rs
@@ -74,7 +74,7 @@ fn test(env: TestEnv) {
         .await
         .unwrap();
 
-        topology_snapshot
+        let topology_snapshot = topology_snapshot
             .block_for_newer_registry_version()
             .await
             .unwrap();