Skip to content

Commit

Permalink
[Core] Fix label validation for Kubernetes, add unit tests (skypilot-…
Browse files Browse the repository at this point in the history
…org#3505)

* Fix Kubernetes label validation and add unit tests

* lint

* lint

* nit
  • Loading branch information
romilbhardwaj authored May 2, 2024
1 parent f256b53 commit d27e0ff
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 3 deletions.
12 changes: 9 additions & 3 deletions sky/clouds/kubernetes.py
Original file line number Diff line number Diff line change
Expand Up @@ -403,11 +403,17 @@ def get_current_user_identity(cls) -> Optional[List[str]]:
@classmethod
def is_label_valid(cls, label_key: str,
label_value: str) -> Tuple[bool, Optional[str]]:
# Kubernetes labels can be of the format <domain>/<key>: <value>
key_regex = re.compile(
r'^(?:[a-z0-9]([-a-z0-9]*[a-z0-9])?\/)?[a-z0-9]([-a-z0-9_.]{0,61}'
r'[a-z0-9])?$')
# Look-ahead to ensure proper domain formatting up to a slash
r'^(?:(?=[a-z0-9]([-a-z0-9.]*[a-z0-9])?\/)'
# Match domain: starts and ends with alphanum up to 253 chars
# including a slash in the domain.
r'[a-z0-9]([-a-z0-9.]{0,251}[a-z0-9])?\/)?'
# Match key: starts and ends with alphanum, upto to 63 chars.
r'[a-z0-9]([-a-z0-9_.]{0,61}[a-z0-9])?$')
value_regex = re.compile(
r'^[a-zA-Z0-9]([-a-zA-Z0-9_.]{0,61}[a-zA-Z0-9])?$')
r'^([a-zA-Z0-9]([-a-zA-Z0-9_.]{0,61}[a-zA-Z0-9])?)?$')
key_valid = bool(key_regex.match(label_key))
value_valid = bool(value_regex.match(label_value))
error_msg = None
Expand Down
75 changes: 75 additions & 0 deletions tests/unit_tests/test_resources.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,23 @@
from typing import Dict
from unittest.mock import Mock

import pytest

from sky import clouds
from sky.resources import Resources

GLOBAL_VALID_LABELS = {
'plaintext': 'plainvalue',
'numbers123': '123',
}

GLOBAL_INVALID_LABELS = {
'l' * 129: 'value', # Long key
'key': 'v' * 257, # Long value
'spaces in label': 'spaces in value',
'': 'emptykey',
}


def test_get_reservations_available_resources():
mock = Mock()
Expand All @@ -11,3 +27,62 @@ def test_get_reservations_available_resources():
r.get_reservations_available_resources()
mock.get_reservations_available_resources.assert_called_once_with(
"instance_type", "region", "zone", set())


def _run_label_test(allowed_labels: Dict[str, str],
invalid_labels: Dict[str, str], cloud: clouds.Cloud):
"""Run a test for labels with the given allowed and invalid labels."""
r_allowed = Resources(cloud=cloud, labels=allowed_labels) # Should pass
assert r_allowed.labels == allowed_labels, ('Allowed labels '
'should be the same')

# Check for each invalid label
for invalid_label, value in invalid_labels.items():
l = {invalid_label: value}
with pytest.raises(ValueError):
_ = Resources(cloud=cloud, labels=l)
assert False, (f'Resources were initialized with '
f'invalid label {invalid_label}={value}')


def test_gcp_labels_resources():
allowed_labels = {
**GLOBAL_VALID_LABELS,
}
invalid_labels = {
**GLOBAL_INVALID_LABELS,
'domain/key': 'value',
'1numericstart': 'value',
}
cloud = clouds.GCP()
_run_label_test(allowed_labels, invalid_labels, cloud)


def test_aws_labels_resources():
allowed_labels = {
**GLOBAL_VALID_LABELS,
}
invalid_labels = {
**GLOBAL_INVALID_LABELS,
'aws:cannotstartwithaws': 'value',
}
cloud = clouds.AWS()
_run_label_test(allowed_labels, invalid_labels, cloud)


def test_kubernetes_labels_resources():
allowed_labels = {
**GLOBAL_VALID_LABELS,
'kueue.x-k8s.io/queue-name': 'queue',
'a' * 253 + '/' + 'k' * 63: 'v' *
63, # upto 253 chars in domain, 63 in key
'mylabel': '', # empty label values are allowed by k8s
'1numericstart': 'value', # numeric start is allowed by k8s
}
invalid_labels = {
**GLOBAL_INVALID_LABELS,
'a' * 254 + '/' + 'k' * 64: 'v' *
63, # exceed 253 chars in domain, 63 in key
}
cloud = clouds.Kubernetes()
_run_label_test(allowed_labels, invalid_labels, cloud)

0 comments on commit d27e0ff

Please sign in to comment.