[Merged by Bors] - Feature: pod builder affinity and node selection

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Added
+
+- Extended the `PodBuilder` with `pod_affinity`, `pod_anti_affinity`, `node_selector` and their `*_opt` variants ([#520])
+
+[#520]: https://github.com/stackabletech/operator-rs/pull/520
+
 ## [0.29.0] - 2022-12-16
 
 ### Added

src/builder/pod/mod.rs

@@ -6,20 +6,18 @@ use crate::builder::meta::ObjectMetaBuilder;
 use crate::commons::product_image_selection::ResolvedProductImage;
 use crate::error::{Error, OperatorResult};
 
+use super::{ListenerOperatorVolumeSourceBuilder, ListenerReference};
 use k8s_openapi::{
     api::core::v1::{
         Affinity, Container, LocalObjectReference, NodeAffinity, NodeSelector,
-        NodeSelectorRequirement, NodeSelectorTerm, Pod, PodAffinity, PodCondition,
+        NodeSelectorRequirement, NodeSelectorTerm, Pod, PodAffinity, PodAntiAffinity, PodCondition,
         PodSecurityContext, PodSpec, PodStatus, PodTemplateSpec, Toleration, Volume,
     },
     apimachinery::pkg::apis::meta::v1::{LabelSelector, LabelSelectorRequirement, ObjectMeta},
 };
 use std::collections::BTreeMap;
 
-use super::{ListenerOperatorVolumeSourceBuilder, ListenerReference};
-
-/// A builder to build [`Pod`] objects.
-///
+/// A builder to build [`Pod`] or [`PodTemplateSpec`] objects.
 #[derive(Clone, Default)]
 pub struct PodBuilder {
     containers: Vec<Container>,
@@ -29,6 +27,7 @@ pub struct PodBuilder {
     node_name: Option<String>,
     node_selector: Option<LabelSelector>,
     pod_affinity: Option<PodAffinity>,
+    pod_anti_affinity: Option<PodAntiAffinity>,
     status: Option<PodStatus>,
     security_context: Option<PodSecurityContext>,
     tolerations: Option<Vec<Toleration>>,
@@ -87,11 +86,31 @@ impl PodBuilder {
         self
     }
 
+    pub fn pod_affinity_opt(&mut self, affinity: Option<PodAffinity>) -> &mut Self {
+        self.pod_affinity = affinity;
+        self
+    }
+
+    pub fn pod_anti_affinity(&mut self, anti_affinity: PodAntiAffinity) -> &mut Self {
+        self.pod_anti_affinity = Some(anti_affinity);
+        self
+    }
+
+    pub fn pod_anti_affinity_opt(&mut self, anti_affinity: Option<PodAntiAffinity>) -> &mut Self {
+        self.pod_anti_affinity = anti_affinity;
+        self
+    }
+
     pub fn node_selector(&mut self, node_selector: LabelSelector) -> &mut Self {
         self.node_selector = Some(node_selector);
         self
     }
 
+    pub fn node_selector_opt(&mut self, node_selector: Option<LabelSelector>) -> &mut Self {
+        self.node_selector = node_selector;
+        self
+    }
+
     pub fn phase(&mut self, phase: &str) -> &mut Self {
         let mut status = self.status.get_or_insert_with(PodStatus::default);
         status.phase = Some(phase.to_string());
@@ -354,18 +373,11 @@ impl PodBuilder {
             init_containers: self.init_containers.clone(),
             node_name: self.node_name.clone(),
             node_selector: node_selector_labels,
-            affinity: node_affinity
-                .map(|node_affinity| Affinity {
-                    node_affinity: Some(node_affinity),
-                    pod_affinity: self.pod_affinity.clone(),
-                    ..Affinity::default()
-                })
-                .or_else(|| {
-                    Some(Affinity {
-                        pod_affinity: self.pod_affinity.clone(),
-                        ..Affinity::default()
-                    })
-                }),
+            affinity: Some(Affinity {
+                node_affinity,
+                pod_affinity: self.pod_affinity.clone(),
+                pod_anti_affinity: self.pod_anti_affinity.clone(),
+            }),
             security_context: self.security_context.clone(),
             tolerations: self.tolerations.clone(),
             volumes: self.volumes.clone(),
@@ -406,33 +418,71 @@ impl PodBuilder {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
     use crate::builder::{
         meta::ObjectMetaBuilder,
-        pod::{container::ContainerBuilder, volume::VolumeBuilder, PodBuilder},
+        pod::{container::ContainerBuilder, volume::VolumeBuilder},
     };
     use k8s_openapi::{
         api::core::v1::{LocalObjectReference, PodAffinity, PodAffinityTerm},
         apimachinery::pkg::apis::meta::v1::{LabelSelector, LabelSelectorRequirement},
     };
+    use rstest::*;
 
-    #[test]
-    fn test_pod_builder() {
-        let container = ContainerBuilder::new("containername")
+    // A simple [`Container`] with a name and image.
+    #[fixture]
+    fn dummy_container() -> Container {
+        ContainerBuilder::new("container")
             .expect("ContainerBuilder not created")
-            .image("stackable/zookeeper:2.4.14")
-            .command(vec!["zk-server-start.sh".to_string()])
-            .args(vec!["stackable/conf/zk.properties".to_string()])
-            .add_volume_mount("zk-worker-1", "conf/")
-            .build();
+            .image("private-company/product:2.4.14")
+            .build()
+    }
 
-        let init_container = ContainerBuilder::new("init-containername")
-            .expect("ContainerBuilder not created")
-            .image("stackable/zookeeper:2.4.14")
-            .command(vec!["wrapper.sh".to_string()])
-            .args(vec!["12345".to_string()])
-            .build();
+    /// A [`PodBuilder`] that already contains the minum setup to build a Pod (name and container).
+    #[fixture]
+    fn pod_builder_with_name_and_container(dummy_container: Container) -> PodBuilder {
+        let mut builder = PodBuilder::new();
+        builder
+            .metadata(ObjectMetaBuilder::new().name("testpod").build())
+            .add_container(dummy_container);
+        builder
+    }
+
+    // A fixture for a node selector to use on a Pod, and the resulting node selector labels and node affinity.
+    #[fixture]
+    fn node_selector1() -> (
+        LabelSelector,
+        Option<BTreeMap<String, String>>,
+        Option<NodeAffinity>,
+    ) {
+        let labels = BTreeMap::from([("key".to_owned(), "value".to_owned())]);
+        let label_selector = LabelSelector {
+            match_expressions: Some(vec![LabelSelectorRequirement {
+                key: "security".to_owned(),
+                operator: "In".to_owned(),
+                values: Some(vec!["S1".to_owned(), "S2".to_owned()]),
+            }]),
+            match_labels: Some(labels.clone()),
+        };
+        let affinity = Some(NodeAffinity {
+            required_during_scheduling_ignored_during_execution: Some(NodeSelector {
+                node_selector_terms: vec![NodeSelectorTerm {
+                    match_expressions: Some(vec![NodeSelectorRequirement {
+                        key: "security".to_owned(),
+                        operator: "In".to_owned(),
+                        values: Some(vec!["S1".to_owned(), "S2".to_owned()]),
+                    }]),
+                    ..Default::default()
+                }],
+            }),
+            ..Default::default()
+        });
+        (label_selector, Some(labels), affinity)
+    }
 
-        let pod_affinity = PodAffinity {
+    #[fixture]
+    fn pod_affinity() -> PodAffinity {
+        PodAffinity {
             preferred_during_scheduling_ignored_during_execution: None,
             required_during_scheduling_ignored_during_execution: Some(vec![PodAffinityTerm {
                 label_selector: Some(LabelSelector {
@@ -446,12 +496,39 @@ mod tests {
                 topology_key: "topology.kubernetes.io/zone".to_string(),
                 ..Default::default()
             }]),
-        };
+        }
+    }
+
+    #[fixture]
+    fn pod_anti_affinity(pod_affinity: PodAffinity) -> PodAntiAffinity {
+        PodAntiAffinity {
+            preferred_during_scheduling_ignored_during_execution: None,
+            required_during_scheduling_ignored_during_execution: pod_affinity
+                .required_during_scheduling_ignored_during_execution,
+        }
+    }
+
+    #[rstest]
+    fn test_pod_builder_pod_name() {
+        let pod = PodBuilder::new()
+            .metadata_builder(|builder| builder.name("foo"))
+            .build()
+            .unwrap();
+
+        assert_eq!(pod.metadata.name.unwrap(), "foo");
+    }
+
+    #[rstest]
+    fn test_pod_builder(pod_affinity: PodAffinity, dummy_container: Container) {
+        let init_container = ContainerBuilder::new("init-containername")
+            .expect("ContainerBuilder not created")
+            .image("stackable/zookeeper:2.4.14")
+            .build();
 
         let pod = PodBuilder::new()
             .pod_affinity(pod_affinity.clone())
             .metadata(ObjectMetaBuilder::new().name("testpod").build())
-            .add_container(container)
+            .add_container(dummy_container)
             .add_init_container(init_container)
             .node_name("worker-1.stackable.demo")
             .add_volume(
@@ -486,25 +563,11 @@ mod tests {
                 .and_then(|volume| volume.config_map.as_ref()?.name.clone())),
             Some("configmap".to_string())
         );
-
-        let pod = PodBuilder::new()
-            .metadata_builder(|builder| builder.name("foo"))
-            .build()
-            .unwrap();
-
-        assert_eq!(pod.metadata.name.unwrap(), "foo");
     }
 
-    #[test]
-    fn test_pod_builder_image_pull_secrets() {
-        let container = ContainerBuilder::new("container")
-            .expect("ContainerBuilder not created")
-            .image("private-comapany/product:2.4.14")
-            .build();
-
-        let pod = PodBuilder::new()
-            .metadata(ObjectMetaBuilder::new().name("testpod").build())
-            .add_container(container)
+    #[rstest]
+    fn test_pod_builder_image_pull_secrets(mut pod_builder_with_name_and_container: PodBuilder) {
+        let pod = pod_builder_with_name_and_container
             .image_pull_secrets(vec!["company-registry-secret".to_string()].into_iter())
             .build()
             .unwrap();
@@ -516,4 +579,88 @@ mod tests {
             }]
         );
     }
+
+    /// Test if setting a node selector generates the correct node selector labels and node affinity on the Pod.
+    #[rstest]
+    fn test_pod_builder_node_selector(
+        mut pod_builder_with_name_and_container: PodBuilder,
+        node_selector1: (
+            LabelSelector,
+            Option<BTreeMap<String, String>>,
+            Option<NodeAffinity>,
+        ),
+    ) {
+        // destruct fixture
+        let (node_selector, expected_labels, expected_affinity) = node_selector1;
+        // first test with the normal node_selector function
+        let pod = pod_builder_with_name_and_container
+            .clone()
+            .node_selector(node_selector.clone())
+            .build()
+            .unwrap();
+
+        let spec = pod.spec.unwrap();
+        assert_eq!(spec.node_selector, expected_labels);
+        assert_eq!(spec.affinity.unwrap().node_affinity, expected_affinity);
+
+        // test the node_selector_opt function
+        let pod = pod_builder_with_name_and_container
+            .node_selector_opt(Some(node_selector))
+            .build()
+            .unwrap();
+
+        // asserts
+        let spec = pod.spec.unwrap();
+        assert_eq!(spec.node_selector, expected_labels);
+        assert_eq!(spec.affinity.unwrap().node_affinity, expected_affinity);
+    }
+
+    /// Test if setting a node selector generates the correct node selector labels and node affinity on the Pod,
+    /// while keeping the manually set Pod affinities. Since they are mangled together, it makes sense to make sure that
+    /// one is not replacing the other
+    #[rstest]
+    fn test_pod_builder_node_selector_and_affinity(
+        mut pod_builder_with_name_and_container: PodBuilder,
+        node_selector1: (
+            LabelSelector,
+            Option<BTreeMap<String, String>>,
+            Option<NodeAffinity>,
+        ),
+        pod_affinity: PodAffinity,
+        pod_anti_affinity: PodAntiAffinity,
+    ) {
+        // destruct fixture
+        let (node_selector, expected_labels, expected_affinity) = node_selector1;
+        // first test with the normal functions
+        let pod = pod_builder_with_name_and_container
+            .clone()
+            .node_selector(node_selector.clone())
+            .pod_affinity(pod_affinity.clone())
+            .pod_anti_affinity(pod_anti_affinity.clone())
+            .build()
+            .unwrap();
+
+        let spec = pod.spec.unwrap();
+        assert_eq!(spec.node_selector, expected_labels);
+        let affinity = spec.affinity.unwrap();
+        assert_eq!(affinity.node_affinity, expected_affinity);
+        assert_eq!(affinity.pod_affinity, Some(pod_affinity.clone()));
+        assert_eq!(affinity.pod_anti_affinity, Some(pod_anti_affinity.clone()));
+
+        // test the *_opt functions
+        let pod = pod_builder_with_name_and_container
+            .node_selector_opt(Some(node_selector))
+            .pod_affinity_opt(Some(pod_affinity.clone()))
+            .pod_anti_affinity_opt(Some(pod_anti_affinity.clone()))
+            .build()
+            .unwrap();
+
+        // asserts
+        let spec = pod.spec.unwrap();
+        assert_eq!(spec.node_selector, expected_labels);
+        let affinity = spec.affinity.unwrap();
+        assert_eq!(affinity.node_affinity, expected_affinity);
+        assert_eq!(affinity.pod_affinity, Some(pod_affinity));
+        assert_eq!(affinity.pod_anti_affinity, Some(pod_anti_affinity));
+    }
 }