def get_replica(self, environment):
main_container = k8s_schemas.V1Container(name="main")
sidecar_containers = [k8s_schemas.V1Container(name="sidecar")]
init_containers = [k8s_schemas.V1Container(name="init")]
replica = ReplicaSpec(
volumes=[],
init_containers=init_containers,
sidecar_containers=sidecar_containers,
main_container=main_container,
labels=environment.labels,
environment=environment,
num_replicas=12,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name="foo",
volumes=[],
environment=environment,
labels=environment.labels,
)
replica_template = {
"replicas":
replica.num_replicas,
"restartPolicy":
pod_spec.restart_policy,
"template":
get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
}
return replica, replica_template
def get_kf_replicas_template(
namespace: str,
resource_name: str,
replica_name: str,
replica: Optional[ReplicaSpec],
labels: Dict[str, str],
annotations: Dict[str, str],
template_spec: Dict,
):
if not replica:
return
metadata, pod_spec = get_pod_spec(
namespace=namespace,
main_container=replica.main_container,
sidecar_containers=replica.sidecar_containers,
init_containers=replica.init_containers,
resource_name=resource_name,
volumes=replica.volumes,
environment=replica.environment,
labels=labels,
annotations=annotations,
)
template_spec[replica_name] = {
"replicas": replica.num_replicas,
"restartPolicy": pod_spec.restart_policy or "Never",
"template": get_pod_template_spec(metadata=metadata,
pod_spec=pod_spec),
}
def test_get_service_custom_resource(self):
main_container = k8s_schemas.V1Container(name="main")
sidecar_containers = [k8s_schemas.V1Container(name="sidecar")]
init_containers = [k8s_schemas.V1Container(name="init")]
termination = V1Termination(timeout=10)
environment = V1Environment(
labels={"foo": "bar"},
annotations={"foo": "bar"},
node_selector={"foo": "bar"},
node_name="foo",
restart_policy="never",
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name="foo",
volumes=[],
environment=environment,
labels=environment.labels,
)
custom_object = {
"serviceSpec": {
"template":
get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
},
"termination": {
"activeDeadlineSeconds": termination.timeout
},
"collectLogs": True,
"syncStatuses": True,
"notifications": [],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels={"foo": "bar"},
custom_object=custom_object,
)
crd = get_service_custom_resource(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
volumes=[],
termination=termination,
environment=environment,
labels=environment.labels,
collect_logs=True,
sync_statuses=True,
notifications=None,
ports=[],
)
assert crd == expected_crd
def test_get_service_custom_resource_missing_keys(self):
main_container = k8s_schemas.V1Container(name="main")
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
resource_name="foo",
volumes=[],
environment=None,
labels=None,
annotations=None,
)
notifications = [
V1Notification(connections=["test"], trigger=V1Statuses.DONE)
]
custom_object = {
"template": get_pod_template_spec(metadata=metadata,
pod_spec=pod_spec),
"ports": [12, 121, 12],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels=None,
annotations=None,
custom_object={
"serviceSpec": custom_object,
"collectLogs": False,
"syncStatuses": False,
"notifications": [n.to_operator() for n in notifications],
},
)
crd = get_service_custom_resource(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
volumes=[],
termination=None,
collect_logs=None,
sync_statuses=None,
notifications=notifications,
environment=None,
labels=None,
annotations=None,
ports=[12, 121, 12],
)
assert crd == expected_crd
def get_service_custom_resource(
resource_name: str,
namespace: str,
main_container: k8s_schemas.V1Container,
sidecar_containers: Optional[List[k8s_schemas.V1Container]],
init_containers: Optional[List[k8s_schemas.V1Container]],
volumes: List[k8s_schemas.V1Volume],
termination: V1Termination,
collect_logs: bool,
sync_statuses: bool,
notifications: List[V1Notification],
environment: V1Environment,
ports: List[int],
labels: Dict[str, str],
annotations: Dict[str, str],
) -> Dict:
metadata, pod_spec = get_pod_spec(
namespace=namespace,
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name=resource_name,
volumes=volumes,
environment=environment,
labels=labels,
)
template_spec = {
"template": get_pod_template_spec(metadata=metadata, pod_spec=pod_spec)
}
if ports:
template_spec["ports"] = ports
custom_object = {"serviceSpec": template_spec}
custom_object = set_termination(
custom_object=custom_object, termination=termination
)
custom_object = set_collect_logs(
custom_object=custom_object, collect_logs=collect_logs
)
custom_object = set_sync_statuses(
custom_object=custom_object, sync_statuses=sync_statuses
)
custom_object = set_notify(custom_object=custom_object, notifications=notifications)
return get_operation_custom_object(
namespace=namespace,
resource_name=resource_name,
labels=labels,
annotations=annotations,
custom_object=custom_object,
)
def test_get_pod_spec(self):
init_container = k8s_schemas.V1Container(name="init")
main_container = k8s_schemas.V1Container(name="main")
sidecar_container = k8s_schemas.V1Container(name="sidecar")
volumes = [k8s_schemas.V1Volume(name="vol")]
labels = {"key": "labels"}
annotations = {"key": "annotations"}
node_selector = {"key": "selector"}
affinity = [{"key": "affinity"}]
tolerations = {"key": "tolerations"}
security_context = {"uid": 222, "gid": 222}
restart_policy = "never"
with self.assertRaises(PolypodException):
get_pod_spec(
namespace="default",
main_container=None,
sidecar_containers=None,
init_containers=None,
resource_name="foo",
volumes=None,
environment=V1Environment(),
labels={},
annotations={},
)
environment = V1Environment(
service_account_name="sa",
labels=labels,
annotations=annotations,
node_selector=node_selector,
affinity=affinity,
tolerations=tolerations,
security_context=security_context,
image_pull_secrets=[],
restart_policy=restart_policy,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=None,
init_containers=None,
volumes=None,
environment=environment,
labels=environment.labels,
annotations=environment.annotations,
)
assert metadata.name == "foo"
assert metadata.labels == labels
assert metadata.namespace == "default"
assert metadata.annotations == annotations
assert isinstance(pod_spec, k8s_schemas.V1PodSpec)
assert pod_spec.security_context == security_context
assert pod_spec.restart_policy == "never"
assert pod_spec.service_account_name == "sa"
assert pod_spec.init_containers == []
assert pod_spec.containers == [main_container]
assert pod_spec.volumes is None
assert pod_spec.node_selector == node_selector
assert pod_spec.tolerations == tolerations
assert pod_spec.affinity == affinity
environment = V1Environment(
service_account_name="sa",
labels=labels,
annotations=annotations,
node_selector=node_selector,
affinity=affinity,
tolerations=tolerations,
security_context=security_context,
image_pull_secrets=[],
restart_policy=restart_policy,
)
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=[sidecar_container],
init_containers=[init_container],
resource_name="foo",
volumes=volumes,
environment=environment,
labels={},
annotations={},
)
assert pod_spec.init_containers == [init_container]
assert pod_spec.containers == [main_container, sidecar_container]
assert pod_spec.volumes == volumes
assert metadata.annotations == {}
def test_get_job_custom_resource(self):
main_container = k8s_schemas.V1Container(name="main")
sidecar_containers = [k8s_schemas.V1Container(name="sidecar")]
init_containers = [k8s_schemas.V1Container(name="init")]
termination = V1Termination(max_retries=5, ttl=10, timeout=10)
environment = V1Environment(
labels={"foo": "bar"},
annotations={"foo": "bar"},
node_selector={"foo": "bar"},
node_name="foo",
restart_policy="never",
)
notifications = [
V1Notification(connections=["test"], trigger=V1NotificationTrigger.DONE)
]
metadata, pod_spec = get_pod_spec(
namespace="default",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
resource_name="foo",
volumes=[],
environment=environment,
labels=environment.labels,
)
custom_object = {
"batchJobSpec": {
"template": get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
},
"termination": {
"backoffLimit": termination.max_retries,
"activeDeadlineSeconds": termination.timeout,
"ttlSecondsAfterFinished": termination.ttl,
},
"collectLogs": True,
"syncStatuses": True,
"notifications": [n.to_operator() for n in notifications],
}
expected_crd = get_custom_object(
namespace="default",
resource_name="foo",
kind="Operation",
api_version="core.polyaxon.com/v1",
labels={"foo": "bar"},
custom_object=custom_object,
annotations={"foo": "long-foo-bar" * 300},
)
crd = get_job_custom_resource(
namespace="default",
resource_name="foo",
main_container=main_container,
sidecar_containers=sidecar_containers,
init_containers=init_containers,
volumes=[],
termination=termination,
environment=environment,
collect_logs=True,
sync_statuses=True,
notifications=notifications,
labels=environment.labels,
annotations={"foo": "long-foo-bar" * 300},
)
assert crd == expected_crd
