fix unit tests

Signed-off-by: Bugen Zhao <[email protected]>

src/common/src/hash/consistent_hash/compat.rs

@@ -38,6 +38,13 @@ impl VnodeCount {
         VnodeCount::Set(NonZeroUsize::new(v).unwrap())
     }
 
+    /// Creates a `VnodeCount` set to the value for testing.
+    ///
+    /// Equivalent to `VnodeCount::set(VirtualNode::COUNT_FOR_TEST)`.
+    pub fn for_test() -> Self {
+        Self::set(VirtualNode::COUNT_FOR_TEST)
+    }
+
     /// Converts to protobuf representation for `maybe_vnode_count`.
     pub fn to_protobuf(self) -> Option<u32> {
         match self {

src/frontend/src/handler/alter_table_column.rs

@@ -19,6 +19,7 @@ use anyhow::{anyhow, Context};
 use itertools::Itertools;
 use pgwire::pg_response::{PgResponse, StatementType};
 use risingwave_common::catalog::ColumnCatalog;
+use risingwave_common::hash::VnodeCount;
 use risingwave_common::types::DataType;
 use risingwave_common::util::column_index_mapping::ColIndexMapping;
 use risingwave_common::{bail, bail_not_implemented};
@@ -244,7 +245,7 @@ pub async fn get_replace_table_plan(
     // Set some fields ourselves so that the meta service does not need to maintain them.
     let mut table = table;
     table.incoming_sinks = incoming_sink_ids.iter().copied().collect();
-    table.maybe_vnode_count = Some(original_catalog.vnode_count() as _);
+    table.maybe_vnode_count = VnodeCount::set(original_catalog.vnode_count()).to_protobuf();
 
     Ok((source, table, graph, col_index_mapping, job_type))
 }

src/frontend/src/test_utils.rs

@@ -28,6 +28,7 @@ use risingwave_common::catalog::{
     FunctionId, IndexId, TableId, DEFAULT_DATABASE_NAME, DEFAULT_SCHEMA_NAME, DEFAULT_SUPER_USER,
     DEFAULT_SUPER_USER_ID, NON_RESERVED_USER_ID, PG_CATALOG_SCHEMA_NAME, RW_CATALOG_SCHEMA_NAME,
 };
+use risingwave_common::hash::{VirtualNode, VnodeCount, VnodeCountCompat};
 use risingwave_common::session_config::SessionConfig;
 use risingwave_common::system_param::reader::SystemParamsReader;
 use risingwave_common::util::cluster_limit::ClusterLimit;
@@ -281,6 +282,7 @@ impl CatalogWriter for MockCatalogWriter {
     ) -> Result<()> {
         table.id = self.gen_id();
         table.stream_job_status = PbStreamJobStatus::Created as _;
+        table.maybe_vnode_count = VnodeCount::for_test().to_protobuf();
         self.catalog.write().create_table(&table);
         self.add_table_or_source_id(table.id, table.schema_id, table.database_id);
         self.hummock_snapshot_manager
@@ -320,6 +322,7 @@ impl CatalogWriter for MockCatalogWriter {
         _job_type: TableJobType,
     ) -> Result<()> {
         table.stream_job_status = PbStreamJobStatus::Created as _;
+        assert_eq!(table.vnode_count(), VirtualNode::COUNT_FOR_TEST);
         self.catalog.write().update_table(&table);
         Ok(())
     }
@@ -353,6 +356,7 @@ impl CatalogWriter for MockCatalogWriter {
     ) -> Result<()> {
         index_table.id = self.gen_id();
         index_table.stream_job_status = PbStreamJobStatus::Created as _;
+        index_table.maybe_vnode_count = VnodeCount::for_test().to_protobuf();
         self.catalog.write().create_table(&index_table);
         self.add_table_or_index_id(
             index_table.id,

src/meta/src/controller/fragment.rs

@@ -1627,7 +1627,7 @@ mod tests {
                 .values()
                 .flat_map(|m| m.keys().map(|x| *x as _))
                 .collect(),
-            maybe_vnode_count: VnodeCount::set(VirtualNode::COUNT_FOR_TEST).to_protobuf(),
+            maybe_vnode_count: VnodeCount::for_test().to_protobuf(),
         };
 
         let pb_actor_status = (0..actor_count)