test: Add tests for TrinoRoutingGroupHeaderRouter

Author: sbernauer

trino-lb/src/routing/python_script.rs

@@ -292,4 +292,22 @@ def targetClusterGroup(query: str, headers: dict[str, str]) -> Optional[str]:
             expected.map(ToOwned::to_owned)
         );
     }
+
+    #[tokio::test]
+    async fn test_invalid_script() {
+        let script = "malformed python :)".to_string();
+        let config = PythonScriptRouterConfig { script };
+
+        let result = PythonScriptRouter::new(&config, HashSet::new());
+        assert!(matches!(result, Err(Error::ParsePythonScript { .. })));
+    }
+
+    #[tokio::test]
+    async fn test_missing_function() {
+        let script = "foo = 42".to_string();
+        let config = PythonScriptRouterConfig { script };
+
+        let result = PythonScriptRouter::new(&config, HashSet::new());
+        assert!(matches!(result, Err(Error::FindPythonFunction { .. })));
+    }
 }

trino-lb/src/routing/trino_routing_group_header.rs

@@ -48,3 +48,83 @@ impl RouterImplementationTrait for TrinoRoutingGroupHeaderRouter {
         None
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use http::{HeaderMap, HeaderName, HeaderValue};
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    #[case(None)]
+    #[case(Some("foo"))]
+    #[case(Some("bar,bak"))]
+    #[tokio::test]
+    async fn test_standard_header(#[case] x_trino_routing_group: Option<&str>) {
+        let config = serde_yaml::from_str("").unwrap();
+        let valid_target_groups = HashSet::from(["foo".to_string(), "bar,bak".to_string()]);
+        let router = TrinoRoutingGroupHeaderRouter::new(&config, valid_target_groups);
+
+        let mut headers = HeaderMap::new();
+        if let Some(x_trino_routing_group) = x_trino_routing_group {
+            headers.insert(
+                HeaderName::from_static("x-trino-routing-group"),
+                HeaderValue::from_str(x_trino_routing_group).unwrap(),
+            );
+        }
+
+        assert_eq!(
+            router.route("", &headers).await.as_deref(),
+            x_trino_routing_group
+        );
+    }
+
+    #[rstest]
+    #[case("x-trino-routing-group", None)]
+    #[case("x-trino-routing-group", Some("foo"))]
+    #[case("x-trino-routing-group", Some("bar,bak"))]
+    #[case("custom-header", None)]
+    #[case("custom-header", Some("foo"))]
+    #[case("custom-header", Some("bar,bak"))]
+    #[tokio::test]
+    async fn test_custom_header(
+        #[case] header_name: String,
+        #[case] x_trino_routing_group: Option<&str>,
+    ) {
+        let config = TrinoRoutingGroupHeaderRouterConfig {
+            header_name: header_name.clone(),
+        };
+        let valid_target_groups = HashSet::from(["foo".to_string(), "bar,bak".to_string()]);
+        let router = TrinoRoutingGroupHeaderRouter::new(&config, valid_target_groups);
+
+        let mut headers = HeaderMap::new();
+        if let Some(x_trino_routing_group) = x_trino_routing_group {
+            headers.insert(
+                HeaderName::from_bytes(header_name.as_bytes()).unwrap(),
+                HeaderValue::from_str(x_trino_routing_group).unwrap(),
+            );
+        }
+
+        assert_eq!(
+            router.route("", &headers).await.as_deref(),
+            x_trino_routing_group
+        );
+    }
+
+    #[tokio::test]
+    async fn test_target_group_does_not_exist() {
+        let config = serde_yaml::from_str("").unwrap();
+        let valid_target_groups = HashSet::from(["foo".to_string()]);
+        let router = TrinoRoutingGroupHeaderRouter::new(&config, valid_target_groups);
+
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            HeaderName::from_static("x-trino-routing-group"),
+            HeaderValue::from_str("does not exist").unwrap(),
+        );
+
+        // Currently we don't raise any error to the user and just ignore this request. This might change in the future.
+        assert_eq!(router.route("", &headers).await, None);
+    }
+}