diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py index 14bbdc2f..e6645f1f 100644 --- a/kubespawner/spawner.py +++ b/kubespawner/spawner.py @@ -781,7 +781,7 @@ def _deprecated_changed(self, change): @validate('image_pull_secrets') def _validate_image_pull_secrets(self, proposal): - if type(proposal['value']) == str: + if isinstance(proposal['value'], str): warnings.warn( """Passing KubeSpawner.image_pull_secrets string values is deprecated since KubeSpawner 0.14.0. The recommended @@ -1262,14 +1262,21 @@ def _validate_image_pull_secrets(self, proposal): """, ) - init_containers = List( + init_containers = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" - List of initialization containers belonging to the pod. + List or dictionary of initialization containers belonging to the pod. - This list will be directly added under `initContainers` in the kubernetes pod spec, - so you should use the same structure. Each item in the dict must a field - of the `V1Container specification `__ + If provided as a list, this list will be directly added under `initContainers` in the kubernetes pod spec. + If provided as a dictionary, the items will be sorted lexicographically by the dictionary keys and + then the sorted values will be added to the `initContainers` key. + + Each item (whether in the list or dictionary values) must be a dictionary which follows the spec at + `V1Container specification `__ One usage is disabling access to metadata service from single-user notebook server with configuration below:: @@ -1285,11 +1292,25 @@ def _validate_image_pull_secrets(self, proposal): } }] + Or as a dictionary:: + + c.KubeSpawner.init_containers = { + "01-iptables": { + "name": "init-iptables", + "image": "", + "command": ["iptables", "-A", "OUTPUT", "-p", "tcp", "--dport", "80", "-d", "169.254.169.254", "-j", "DROP"], + "securityContext": { + "capabilities": { + "add": ["NET_ADMIN"] + } + } + } + } See `the Kubernetes documentation `__ for more info on what init containers are and why you might want to use them! - To user this feature, Kubernetes version must greater than 1.6. + To use this feature, Kubernetes version must be greater than 1.6. """, ) @@ -1342,15 +1363,22 @@ def _validate_image_pull_secrets(self, proposal): """, ) - extra_containers = List( + extra_containers = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" - List of containers belonging to the pod which besides to the container generated for notebook server. + List or dictionary of containers belonging to the pod in addition to + the container generated for the notebook server. - This list will be directly appended under `containers` in the kubernetes pod spec, - so you should use the same structure. Each item in the list is container configuration - which follows spec at https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#container-v1-core + If provided as a list, this list will be directly appended under `containers` in the kubernetes pod spec. + If provided as a dictionary, the items will be sorted lexicographically by the dictionary keys and + then the sorted values will be appended to the `containers` key. + Each item (whether in the list or dictionary values) is container configuration + which follows the spec at https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#container-v1-core One usage is setting crontab in a container to clean sensitive data with configuration below:: @@ -1360,6 +1388,16 @@ def _validate_image_pull_secrets(self, proposal): "command": ["/usr/local/bin/supercronic", "/etc/crontab"] }] + Or as a dictionary:: + + c.KubeSpawner.extra_containers = { + "01-crontab": { + "name": "crontab", + "image": "supercronic", + "command": ["/usr/local/bin/supercronic", "/etc/crontab"] + } + } + `{username}`, `{userid}`, `{servername}`, `{hubnamespace}`, `{unescaped_username}`, and `{unescaped_servername}` will be expanded if found within strings of this configuration. The username and servername @@ -1379,19 +1417,27 @@ def _validate_image_pull_secrets(self, proposal): """, ) - tolerations = List( + tolerations = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" - List of tolerations that are to be assigned to the pod in order to be able to schedule the pod + List or dictionary of tolerations that are to be assigned to the pod in order to be able to schedule the pod on a node with the corresponding taints. See the official Kubernetes documentation for additional details https://kubernetes.io/docs/concepts/configuration/taint-and-toleration/ - Pass this field an array of "Toleration" objects - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#toleration-v1-core + If provided as a list, each item should be a "Toleration" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "Toleration" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. - Example:: + Each "Toleration" object should follow the specification at: + https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#toleration-v1-core - [ + Example as a list:: + + c.KubeSpawner.tolerations = [ { 'key': 'key', 'operator': 'Equal', @@ -1405,83 +1451,160 @@ def _validate_image_pull_secrets(self, proposal): } ] + Example as a dictionary:: + + c.KubeSpawner.tolerations = { + "01-gpu-toleration": { + 'key': 'gpu', + 'operator': 'Equal', + 'value': 'true', + 'effect': 'NoSchedule' + }, + "02-general-toleration": { + 'key': 'key', + 'operator': 'Exists', + 'effect': 'NoSchedule' + } + } + """, ) - node_affinity_preferred = List( + node_affinity_preferred = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of preferred node affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "PreferredSchedulingTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#preferredschedulingterm-v1-core + If provided as a list, each item should be a "PreferredSchedulingTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "PreferredSchedulingTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + Each item should follow the `"PreferredSchedulingTerm" specification + `__. """, ) - node_affinity_required = List( + + node_affinity_required = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of required node affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "NodeSelectorTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#nodeselectorterm-v1-core + If provided as a list, each item should be a "NodeSelectorTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "NodeSelectorTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + Each item should follow the `"NodeSelectorTerm" specification + `__. """, ) - pod_affinity_preferred = List( + + pod_affinity_preferred = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of preferred pod affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "WeightedPodAffinityTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#weightedpodaffinityterm-v1-core + If provided as a list, each item should be a "WeightedPodAffinityTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "WeightedPodAffinityTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + Each item should follow the `"WeightedPodAffinityTerm" specification + `__. """, ) - pod_affinity_required = List( + + pod_affinity_required = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of required pod affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "PodAffinityTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#podaffinityterm-v1-core + If provided as a list, each item should be a "PodAffinityTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "PodAffinityTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + Each item should follow the `"PodAffinityTerm" specification + `__. """, ) - pod_anti_affinity_preferred = List( + + pod_anti_affinity_preferred = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of preferred pod anti-affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "WeightedPodAffinityTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#weightedpodaffinityterm-v1-core + If provided as a list, each item should be a "WeightedPodAffinityTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "WeightedPodAffinityTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + + Each item should follow the `"WeightedPodAffinityTerm" specification + `__. """, ) - pod_anti_affinity_required = List( + + pod_anti_affinity_required = Union( + trait_types=[ + List(), + Dict(), + ], config=True, help=""" + List or dictionary of required pod anti-affinities. + Affinities describe where pods prefer or require to be scheduled, they may prefer or require a node to have a certain label or be in proximity / remoteness to another pod. To learn more visit https://kubernetes.io/docs/concepts/configuration/assign-pod-node/ - Pass this field an array of "PodAffinityTerm" objects.* - * https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#podaffinityterm-v1-core + If provided as a list, each item should be a "PodAffinityTerm" object. + If provided as a dictionary, the keys can be any descriptive name and the values should be "PodAffinityTerm" objects. + The items will be sorted lexicographically by the dictionary keys and the sorted values will be added to the pod spec. + + Each item should follow the `"PodAffinityTerm" specification + `__. """, ) @@ -2077,20 +2200,34 @@ async def get_pod_manifest(self): extra_resource_limits=self.extra_resource_limits, extra_resource_guarantees=self.extra_resource_guarantees, lifecycle_hooks=self.lifecycle_hooks, - init_containers=self._expand_all(self.init_containers), + init_containers=self._expand_all( + self._sorted_dict_values(self.init_containers) + ), service_account=self._expand_all(self.service_account), automount_service_account_token=self.automount_service_account_token, extra_container_config=self.extra_container_config, extra_pod_config=self._expand_all(self.extra_pod_config), - extra_containers=self._expand_all(self.extra_containers), + extra_containers=self._expand_all( + self._sorted_dict_values(self.extra_containers) + ), scheduler_name=self.scheduler_name, - tolerations=self.tolerations, - node_affinity_preferred=self.node_affinity_preferred, - node_affinity_required=self.node_affinity_required, - pod_affinity_preferred=self.pod_affinity_preferred, - pod_affinity_required=self.pod_affinity_required, - pod_anti_affinity_preferred=self.pod_anti_affinity_preferred, - pod_anti_affinity_required=self.pod_anti_affinity_required, + tolerations=self._sorted_dict_values(self.tolerations), + node_affinity_preferred=self._sorted_dict_values( + self.node_affinity_preferred + ), + node_affinity_required=self._sorted_dict_values( + self.node_affinity_required + ), + pod_affinity_preferred=self._sorted_dict_values( + self.pod_affinity_preferred + ), + pod_affinity_required=self._sorted_dict_values(self.pod_affinity_required), + pod_anti_affinity_preferred=self._sorted_dict_values( + self.pod_anti_affinity_preferred + ), + pod_anti_affinity_required=self._sorted_dict_values( + self.pod_anti_affinity_required + ), priority_class_name=self.priority_class_name, ssl_secret_name=self.secret_name if self.internal_ssl else None, ssl_secret_mount_path=self.secret_mount_path, diff --git a/tests/test_spawner.py b/tests/test_spawner.py index 7c3c7236..68e8ee60 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -901,6 +901,41 @@ async def test_expansion_hyphens(): ] +async def test_init_containers_as_dict(): + """ + Test that the init_containers config option can be a dictionary of lists + and the values are sorted by key before being added to the pod manifest. + """ + c = Config() + + c.KubeSpawner.init_containers = { + '02-beta': { + 'name': 'mock_name_2', + 'image': 'mock_image_2', + 'command': ['mock_command_2'], + }, + '01-alpha': { + 'name': 'mock_name_1', + 'image': 'mock_image_1', + 'command': ['mock_command_1'], + }, + } + + spawner = KubeSpawner(config=c, _mock=True) + + manifest = await spawner.get_pod_manifest() + assert isinstance(manifest, V1Pod) + init_containers = manifest.spec.init_containers + assert len(init_containers) == 2 + for container in init_containers: + assert isinstance(container, V1Container) + + assert init_containers[0].name == 'mock_name_1' + assert init_containers[1].name == 'mock_name_2' + assert init_containers[0].image == 'mock_image_1' + assert init_containers[1].image == 'mock_image_2' + + _test_profiles = [ { 'display_name': 'Training Env - Python', @@ -1914,3 +1949,606 @@ async def test_volume_list(): ) assert manifest.spec.volumes[1].name == 'volumes-beta' assert manifest.spec.volumes[1].persistent_volume_claim["claimName"] == 'beta-claim' + + +async def test_extra_containers_dictionary(): + """ + Test that extra_containers can be a dictionary of dictionaries. + The output list should be lexicographically sorted by key. + """ + c = Config() + + c.KubeSpawner.extra_containers = { + "02-group-beta": { + 'name': 'extra-containers-beta', + 'image': 'busybox', + }, + "01-group-alpha": { + 'name': 'extra-containers-alpha', + 'image': 'busybox', + }, + } + + spawner = KubeSpawner(config=c, _mock=True) + + manifest = await spawner.get_pod_manifest() + + assert isinstance(manifest.spec.containers, list) + assert ( + len(manifest.spec.containers) == 3 + ) # 1 for the notebook container, 2 for extra_containers + assert manifest.spec.containers[1].name == 'extra-containers-alpha' + assert manifest.spec.containers[2].name == 'extra-containers-beta' + + +async def test_extra_containers_list(): + """ + Test that extra_containers can be a list of dictionaries. + """ + c = Config() + + c.KubeSpawner.extra_containers = [ + { + 'name': 'extra-containers-alpha', + 'image': 'busybox', + }, + { + 'name': 'extra-containers-beta', + 'image': 'busybox', + }, + ] + spawner = KubeSpawner(config=c, _mock=True) + + manifest = await spawner.get_pod_manifest() + + assert isinstance(manifest.spec.containers, list) + assert ( + len(manifest.spec.containers) == 3 + ) # 1 for the notebook container, 2 for extra_containers + assert manifest.spec.containers[1].name == 'extra-containers-alpha' + assert manifest.spec.containers[2].name == 'extra-containers-beta' + + +async def test_tolerations_list(): + """ + Test that tolerations can be a list. + """ + c = Config() + + tolerations = [ + { + 'key': 'key1', + 'operator': 'Equal', + 'value': 'value1', + 'effect': 'NoSchedule', + }, + { + 'key': 'key2', + 'operator': 'Exists', + 'effect': 'NoExecute', + }, + ] + c.KubeSpawner.tolerations = tolerations + spawner = KubeSpawner(config=c, _mock=True) + + manifest = await spawner.get_pod_manifest() + + assert isinstance(manifest.spec.tolerations, list) + assert len(manifest.spec.tolerations) == 2 + assert manifest.spec.tolerations[0].key == 'key1' + assert manifest.spec.tolerations[0].operator == 'Equal' + assert manifest.spec.tolerations[0].value == 'value1' + assert manifest.spec.tolerations[0].effect == 'NoSchedule' + assert manifest.spec.tolerations[1].key == 'key2' + assert manifest.spec.tolerations[1].operator == 'Exists' + assert manifest.spec.tolerations[1].effect == 'NoExecute' + + +async def test_tolerations_dict(): + """ + Test that tolerations can be a dictionary. + The output list should be lexicographically sorted by key. + """ + c = Config() + + tolerations = { + '01-group-alpha': { + 'key': 'key1', + 'operator': 'Equal', + 'value': 'value1', + 'effect': 'NoSchedule', + }, + '02-group-beta': { + 'key': 'key2', + 'operator': 'Exists', + 'effect': 'NoExecute', + }, + } + c.KubeSpawner.tolerations = tolerations + spawner = KubeSpawner(config=c, _mock=True) + + manifest = await spawner.get_pod_manifest() + + assert isinstance(manifest.spec.tolerations, list) + assert len(manifest.spec.tolerations) == 2 + assert manifest.spec.tolerations[0].key == 'key1' + assert manifest.spec.tolerations[0].operator == 'Equal' + assert manifest.spec.tolerations[0].value == 'value1' + assert manifest.spec.tolerations[0].effect == 'NoSchedule' + assert manifest.spec.tolerations[1].key == 'key2' + assert manifest.spec.tolerations[1].operator == 'Exists' + assert manifest.spec.tolerations[1].effect == 'NoExecute' + + +async def test_node_affinity_preferred(): + """ + Test that node_affinity_preferred can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + node_affinity_preferred = [ + { + 'weight': 1, + 'preference': { + 'matchExpressions': [ + {'key': 'key1', 'operator': 'In', 'values': ['value1', 'value2']} + ] + }, + }, + { + 'weight': 2, + 'preference': {'matchExpressions': [{'key': 'key2', 'operator': 'Exists'}]}, + }, + ] + c.KubeSpawner.node_affinity_preferred = node_affinity_preferred + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.node_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert isinstance(spec, list) + assert len(spec) == 2 + assert spec[0].weight == 1 + assert spec[0].preference == node_affinity_preferred[0]["preference"] + assert spec[1].weight == 2 + assert spec[1].preference == node_affinity_preferred[1]["preference"] + + node_affinity_preferred_dict = { + "02-group-beta": { + "weight": 2, + "preference": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + }, + "01-group-alpha": { + "weight": 1, + "preference": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + }, + } + c.KubeSpawner.node_affinity_preferred = node_affinity_preferred_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.node_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert len(spec) == 2 + assert spec[0].weight == 1 + assert ( + spec[0].preference + == node_affinity_preferred_dict["01-group-alpha"]["preference"] + ) + assert spec[1].weight == 2 + assert ( + spec[1].preference + == node_affinity_preferred_dict["02-group-beta"]["preference"] + ) + + +async def test_node_affinity_required(): + """ + Test that node_affinity_required can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + node_affinity_required = [ + { + 'matchExpressions': [ + {'key': 'key1', 'operator': 'In', 'values': ['value1', 'value2']} + ] + } + ] + c.KubeSpawner.node_affinity_required = node_affinity_required + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.node_affinity.required_during_scheduling_ignored_during_execution.node_selector_terms + ) + + assert isinstance(spec, list) + assert len(spec) == 1 + assert spec[0].match_expressions == node_affinity_required[0]["matchExpressions"] + + node_affinity_required_dict = { + "02-group-beta": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + "01-group-alpha": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + } + c.KubeSpawner.node_affinity_required = node_affinity_required_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.node_affinity.required_during_scheduling_ignored_during_execution.node_selector_terms + ) + + assert len(spec) == 2 + assert ( + spec[0].match_expressions + == node_affinity_required_dict["01-group-alpha"]["matchExpressions"] + ) + assert ( + spec[1].match_expressions + == node_affinity_required_dict["02-group-beta"]["matchExpressions"] + ) + + +async def test_pod_affinity_preferred(): + """ + Test that pod_affinity_preferred can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + pod_affinity_preferred = [ + { + 'weight': 1, + 'podAffinityTerm': { + 'labelSelector': { + 'matchExpressions': [ + {'key': 'key1', 'operator': 'In', 'values': ['value1']} + ] + }, + 'topologyKey': 'topology.kubernetes.io/zone', + }, + } + ] + c.KubeSpawner.pod_affinity_preferred = pod_affinity_preferred + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert isinstance(spec, list) + assert len(spec) == 1 + assert spec[0].weight == 1 + assert spec[0].pod_affinity_term == pod_affinity_preferred[0]["podAffinityTerm"] + + pod_affinity_preferred_dict = { + "02-group-beta": { + "weight": 2, + "podAffinityTerm": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/zone", + }, + }, + "01-group-alpha": { + "weight": 1, + "podAffinityTerm": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/zone", + }, + }, + } + c.KubeSpawner.pod_affinity_preferred = pod_affinity_preferred_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert len(spec) == 2 + assert spec[0].weight == 1 + assert ( + spec[0].pod_affinity_term + == pod_affinity_preferred_dict["01-group-alpha"]["podAffinityTerm"] + ) + assert spec[1].weight == 2 + assert ( + spec[1].pod_affinity_term + == pod_affinity_preferred_dict["02-group-beta"]["podAffinityTerm"] + ) + + +async def test_pod_affinity_required(): + """ + Test that pod_affinity_required can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + pod_affinity_required = [ + { + 'labelSelector': { + 'matchExpressions': [ + {'key': 'key1', 'operator': 'In', 'values': ['value1']} + ] + }, + 'topologyKey': 'topology.kubernetes.io/zone', + } + ] + c.KubeSpawner.pod_affinity_required = pod_affinity_required + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_affinity.required_during_scheduling_ignored_during_execution + ) + + assert isinstance(spec, list) + assert len(spec) == 1 + assert spec[0].label_selector == pod_affinity_required[0]["labelSelector"] + assert spec[0].topology_key == pod_affinity_required[0]["topologyKey"] + + pod_affinity_required_dict = { + "02-group-beta": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/zone", + }, + "01-group-alpha": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/zone", + }, + } + c.KubeSpawner.pod_affinity_required = pod_affinity_required_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_affinity.required_during_scheduling_ignored_during_execution + ) + + assert len(spec) == 2 + assert ( + spec[0].label_selector + == pod_affinity_required_dict["01-group-alpha"]["labelSelector"] + ) + assert ( + spec[0].topology_key + == pod_affinity_required_dict["01-group-alpha"]["topologyKey"] + ) + assert ( + spec[1].label_selector + == pod_affinity_required_dict["02-group-beta"]["labelSelector"] + ) + assert ( + spec[1].topology_key + == pod_affinity_required_dict["02-group-beta"]["topologyKey"] + ) + + +async def test_pod_anti_affinity_preferred(): + """ + Test that pod_anti_affinity_preferred can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + pod_anti_affinity_preferred = [ + { + 'weight': 1, + 'podAffinityTerm': { + 'labelSelector': { + 'matchExpressions': [ + {'key': 'key1', 'operator': 'In', 'values': ['value1']} + ] + }, + 'topologyKey': 'topology.kubernetes.io/zone', + }, + } + ] + c.KubeSpawner.pod_anti_affinity_preferred = pod_anti_affinity_preferred + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + + spec = ( + manifest.spec.affinity.pod_anti_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert isinstance(spec, list) + assert len(spec) == 1 + assert spec[0].weight == 1 + assert ( + spec[0].pod_affinity_term == pod_anti_affinity_preferred[0]["podAffinityTerm"] + ) + + pod_anti_affinity_preferred_dict = { + "02-group-beta": { + "weight": 2, + "podAffinityTerm": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/zone", + }, + }, + "01-group-alpha": { + "weight": 1, + "podAffinityTerm": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + "topologyKey": "topology.kubernetes.io/region", + }, + }, + } + c.KubeSpawner.pod_anti_affinity_preferred = pod_anti_affinity_preferred_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_anti_affinity.preferred_during_scheduling_ignored_during_execution + ) + + assert len(spec) == 2 + assert spec[0].weight == 1 + assert ( + spec[0].pod_affinity_term + == pod_anti_affinity_preferred_dict["01-group-alpha"]["podAffinityTerm"] + ) + assert spec[1].weight == 2 + assert ( + spec[1].pod_affinity_term + == pod_anti_affinity_preferred_dict["02-group-beta"]["podAffinityTerm"] + ) + + +async def test_pod_anti_affinity_required(): + """ + Test that pod_anti_affinity_required can be a list or dictionary of dictionaries. + The output list should be lexicographically sorted by key when a dictionary is used. + """ + c = Config() + pod_anti_affinity_required = [ + { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + "topologyKey": "failure-domain.beta.kubernetes.io/zone", + } + ] + c.KubeSpawner.pod_anti_affinity_required = pod_anti_affinity_required + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + + spec = ( + manifest.spec.affinity.pod_anti_affinity.required_during_scheduling_ignored_during_execution + ) + + assert spec[0].label_selector == pod_anti_affinity_required[0]["labelSelector"] + assert spec[0].topology_key == pod_anti_affinity_required[0]["topologyKey"] + + pod_anti_affinity_required_dict = { + "02-group-beta": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S2"], + } + ] + }, + "topologyKey": "failure-domain.beta.kubernetes.io/region", + }, + "01-group-alpha": { + "labelSelector": { + "matchExpressions": [ + { + "key": "security", + "operator": "In", + "values": ["S1"], + } + ] + }, + "topologyKey": "failure-domain.beta.kubernetes.io/zone", + }, + } + c.KubeSpawner.pod_anti_affinity_required = pod_anti_affinity_required_dict + spawner = KubeSpawner(config=c, _mock=True) + manifest = await spawner.get_pod_manifest() + spec = ( + manifest.spec.affinity.pod_anti_affinity.required_during_scheduling_ignored_during_execution + ) + + assert len(spec) == 2 + assert ( + spec[0].label_selector + == pod_anti_affinity_required_dict["01-group-alpha"]["labelSelector"] + ) + assert ( + spec[0].topology_key + == pod_anti_affinity_required_dict["01-group-alpha"]["topologyKey"] + ) + assert ( + spec[1].label_selector + == pod_anti_affinity_required_dict["02-group-beta"]["labelSelector"] + ) + assert ( + spec[1].topology_key + == pod_anti_affinity_required_dict["02-group-beta"]["topologyKey"] + )