From ea3e77e5d574fcdb75ad52b47178451843955d82 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Uzarski?= <mikolaj.uzarski@scylladb.com>
Date: Thu, 17 Oct 2024 15:15:15 +0200
Subject: [PATCH] tests: additional unit test cases for lbp config

Added test cases for rack-awareness, and extended dc-awareness
tests by empty and nullptr parameters checks.
---
 scylla-rust-wrapper/src/cluster.rs | 88 ++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/scylla-rust-wrapper/src/cluster.rs b/scylla-rust-wrapper/src/cluster.rs
index 0a9fe53f..14ffbf91 100644
--- a/scylla-rust-wrapper/src/cluster.rs
+++ b/scylla-rust-wrapper/src/cluster.rs
@@ -929,6 +929,49 @@ mod tests {
                     }
                     assert!(!cluster.load_balancing_config.token_awareness_enabled);
                     assert!(cluster.load_balancing_config.latency_awareness_enabled);
+
+                    // set preferred rack+dc
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_rack_aware(
+                            cluster_raw,
+                            "eu-east\0".as_ptr() as *const i8,
+                            "rack1\0".as_ptr() as *const i8,
+                        ),
+                        CassError::CASS_OK
+                    );
+
+                    let node_location_preference =
+                        &cluster.load_balancing_config.node_location_preference;
+                    match node_location_preference {
+                        NodeLocationPreference::DatacenterAndRack {
+                            local_dc,
+                            local_rack,
+                        } => {
+                            assert_eq!(local_dc, "eu-east");
+                            assert_eq!(local_rack, "rack1");
+                        }
+                        _ => panic!("Expected preferred dc and rack"),
+                    }
+
+                    // set back to preferred dc
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_dc_aware(
+                            cluster_raw,
+                            "eu\0".as_ptr() as *const i8,
+                            0,
+                            0
+                        ),
+                        CassError::CASS_OK
+                    );
+
+                    let node_location_preference =
+                        &cluster.load_balancing_config.node_location_preference;
+                    match node_location_preference {
+                        NodeLocationPreference::Datacenter { local_dc } => {
+                            assert_eq!(local_dc, "eu")
+                        }
+                        _ => panic!("Expected preferred dc"),
+                    }
                 }
                 /* Test invalid configurations */
                 {
@@ -951,6 +994,51 @@ mod tests {
                         ),
                         CassError::CASS_ERROR_LIB_BAD_PARAMS
                     );
+
+                    // null pointers
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_dc_aware(cluster_raw, std::ptr::null(), 0, 0),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_rack_aware(
+                            cluster_raw,
+                            "eu\0".as_ptr() as *const i8,
+                            std::ptr::null(),
+                        ),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_rack_aware(
+                            cluster_raw,
+                            std::ptr::null(),
+                            "rack\0".as_ptr() as *const i8,
+                        ),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
+
+                    // empty strings
+                    let empty_str = "\0".as_ptr() as *const i8;
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_dc_aware(cluster_raw, std::ptr::null(), 0, 0),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_rack_aware(
+                            cluster_raw,
+                            "eu\0".as_ptr() as *const i8,
+                            empty_str,
+                        ),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
+                    assert_cass_error_eq!(
+                        cass_cluster_set_load_balance_rack_aware(
+                            cluster_raw,
+                            empty_str,
+                            "rack\0".as_ptr() as *const i8,
+                        ),
+                        CassError::CASS_ERROR_LIB_BAD_PARAMS
+                    );
                 }
             }