Skip to content

Commit

Permalink
Pod Node Affinity
Browse files Browse the repository at this point in the history
  • Loading branch information
jbartosik committed Dec 27, 2024
1 parent 56ffd03 commit 8948b24
Show file tree
Hide file tree
Showing 2 changed files with 300 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,72 @@ func ExtractPod(p *corev1.Pod) *model.Pod {
}
}

if p.Spec.Affinity != nil && p.Spec.Affinity.NodeAffinity != nil {
podModel.NodeAffinity = &model.NodeAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: convertNodeSelector(p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution),
PreferredDuringSchedulingIgnoredDuringExecution: convertPreferredSchedulingTerm(p.Spec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution),
}
}

return &podModel
}

func convertNodeSelector(ns *corev1.NodeSelector) *model.NodeSelector {
if ns == nil {
return nil
}
return &model.NodeSelector{
NodeSelectorTerms: convertNodeSelectorTerms(ns.NodeSelectorTerms),
}
}

func convertPreferredSchedulingTerm(terms []corev1.PreferredSchedulingTerm) []*model.PreferredSchedulingTerm {
if len(terms) == 0 {
return nil
}
var preferredTerms []*model.PreferredSchedulingTerm
for _, term := range terms {
preferredTerms = append(preferredTerms, &model.PreferredSchedulingTerm{
Preference: convertNodeSelectorTerm(term.Preference),
Weight: term.Weight,
})
}
return preferredTerms
}

func convertNodeSelectorTerms(terms []corev1.NodeSelectorTerm) []*model.NodeSelectorTerm {
if len(terms) == 0 {
return nil
}
var nodeSelectorTerms []*model.NodeSelectorTerm
for _, term := range terms {
nodeSelectorTerms = append(nodeSelectorTerms, convertNodeSelectorTerm(term))
}
return nodeSelectorTerms
}

func convertNodeSelectorTerm(term corev1.NodeSelectorTerm) *model.NodeSelectorTerm {
return &model.NodeSelectorTerm{
MatchExpressions: convertNodeSelectorRequirements(term.MatchExpressions),
MatchFields: convertNodeSelectorRequirements(term.MatchFields),
}
}

func convertNodeSelectorRequirements(requirements []corev1.NodeSelectorRequirement) []*model.LabelSelectorRequirement {
if len(requirements) == 0 {
return nil
}
var nodeSelectorRequirements []*model.LabelSelectorRequirement
for _, req := range requirements {
nodeSelectorRequirements = append(nodeSelectorRequirements, &model.LabelSelectorRequirement{
Key: req.Key,
Operator: string(req.Operator),
Values: req.Values,
})
}
return nodeSelectorRequirements
}

// ExtractPodTemplateResourceRequirements extracts resource requirements of containers and initContainers into model.ResourceRequirements
func ExtractPodTemplateResourceRequirements(template corev1.PodTemplateSpec) []*model.ResourceRequirements {
return extractPodResourceRequirements(template.Spec.Containers, template.Spec.InitContainers)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,14 @@ package k8s

import (
"fmt"
"reflect"
"testing"
"time"

model "github.com/DataDog/agent-payload/v5/process"

"github.com/stretchr/testify/assert"
corev1 "k8s.io/api/core/v1"
v1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
Expand Down Expand Up @@ -966,3 +968,238 @@ func TestMapToTags(t *testing.T) {
assert.ElementsMatch(t, []string{"foo:bar", "node-role.kubernetes.io/nodeless"}, tags)
assert.Len(t, tags, 2)
}

func TestConvertNodeSelector(t *testing.T) {
tests := []struct {
name string
input *corev1.NodeSelector
want *model.NodeSelector
}{
{
name: "nil input",
input: nil,
want: nil,
},
{
name: "empty NodeSelector",
input: &corev1.NodeSelector{
NodeSelectorTerms: []corev1.NodeSelectorTerm{},
},
want: &model.NodeSelector{NodeSelectorTerms: nil},
},
{
name: "with MatchExpressions and MatchFields",
input: &corev1.NodeSelector{
NodeSelectorTerms: []corev1.NodeSelectorTerm{
{
MatchExpressions: []corev1.NodeSelectorRequirement{
{Key: "key1", Operator: corev1.NodeSelectorOpIn, Values: []string{"v1", "v2"}},
},
MatchFields: []corev1.NodeSelectorRequirement{
{Key: "field1", Operator: corev1.NodeSelectorOpNotIn, Values: []string{"v3"}},
},
},
},
},
want: &model.NodeSelector{
NodeSelectorTerms: []*model.NodeSelectorTerm{
{
MatchExpressions: []*model.NodeSelectorRequirement{
{Key: "key1", Operator: "In", Values: []string{"v1", "v2"}},
},
MatchFields: []*model.NodeSelectorRequirement{
{Key: "field1", Operator: "NotIn", Values: []string{"v3"}},
},
},
},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := convertNodeSelector(tt.input)
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("convertNodeSelector() = %#v, want %#v", got, tt.want)
}
})
}
}

func TestConvertPreferredSchedulingTerm(t *testing.T) {
tests := []struct {
name string
input []corev1.PreferredSchedulingTerm
want []*model.PreferredSchedulingTerm
}{
{
name: "empty terms",
input: []corev1.PreferredSchedulingTerm{},
want: nil,
},
{
name: "single preferred scheduling term",
input: []corev1.PreferredSchedulingTerm{
{
Preference: corev1.NodeSelectorTerm{
MatchExpressions: []corev1.NodeSelectorRequirement{
{Key: "k", Operator: corev1.NodeSelectorOpExists},
},
},
Weight: 10,
},
},
want: []*model.PreferredSchedulingTerm{
{
Preference: &model.NodeSelectorTerm{
MatchExpressions: []*model.NodeSelectorRequirement{
{Key: "k", Operator: "Exists", Values: nil},
},
MatchFields: nil,
},
Weight: 10,
},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := convertPreferredSchedulingTerm(tt.input)
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("convertPreferredSchedulingTerm() = %#v, want %#v", got, tt.want)
}
})
}
}

func TestConvertNodeSelectorTerms(t *testing.T) {
tests := []struct {
name string
input []corev1.NodeSelectorTerm
want []*model.NodeSelectorTerm
}{
{
name: "empty terms",
input: []corev1.NodeSelectorTerm{},
want: nil,
},
{
name: "multiple NodeSelectorTerms",
input: []corev1.NodeSelectorTerm{
{
MatchExpressions: []corev1.NodeSelectorRequirement{
{Key: "k1", Operator: corev1.NodeSelectorOpIn, Values: []string{"v1"}},
},
},
{
MatchExpressions: []corev1.NodeSelectorRequirement{
{Key: "k2", Operator: corev1.NodeSelectorOpNotIn, Values: []string{"v2"}},
},
},
},
want: []*model.NodeSelectorTerm{
{
MatchExpressions: []*model.NodeSelectorRequirement{
{Key: "k1", Operator: "In", Values: []string{"v1"}},
},
MatchFields: nil,
},
{
MatchExpressions: []*model.NodeSelectorRequirement{
{Key: "k2", Operator: "NotIn", Values: []string{"v2"}},
},
MatchFields: nil,
},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := convertNodeSelectorTerms(tt.input)
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("convertNodeSelectorTerms() = %#v, want %#v", got, tt.want)
}
})
}
}

func TestConvertNodeSelectorTerm(t *testing.T) {
tests := []struct {
name string
input corev1.NodeSelectorTerm
want *model.NodeSelectorTerm
}{
{
name: "empty term",
input: corev1.NodeSelectorTerm{},
want: &model.NodeSelectorTerm{
MatchExpressions: nil,
MatchFields: nil,
},
},
{
name: "with match expressions and fields",
input: corev1.NodeSelectorTerm{
MatchExpressions: []corev1.NodeSelectorRequirement{
{Key: "k1", Operator: corev1.NodeSelectorOpExists},
},
MatchFields: []corev1.NodeSelectorRequirement{
{Key: "f1", Operator: corev1.NodeSelectorOpDoesNotExist},
},
},
want: &model.NodeSelectorTerm{
MatchExpressions: []*model.NodeSelectorRequirement{
{Key: "k1", Operator: "Exists"},
},
MatchFields: []*model.NodeSelectorRequirement{
{Key: "f1", Operator: "DoesNotExist"},
},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := convertNodeSelectorTerm(tt.input)
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("convertNodeSelectorTerm() = %#v, want %#v", got, tt.want)
}
})
}
}

func TestConvertNodeSelectorRequirements(t *testing.T) {
tests := []struct {
name string
input []corev1.NodeSelectorRequirement
want []*model.LabelSelectorRequirement
}{
{
name: "no requirements",
input: []corev1.NodeSelectorRequirement{},
want: nil,
},
{
name: "with multiple requirements",
input: []corev1.NodeSelectorRequirement{
{Key: "k1", Operator: corev1.NodeSelectorOpIn, Values: []string{"v1", "v2"}},
{Key: "k2", Operator: corev1.NodeSelectorOpNotIn, Values: []string{"v3"}},
},
want: []*model.LabelSelectorRequirement{
{Key: "k1", Operator: "In", Values: []string{"v1", "v2"}},
{Key: "k2", Operator: "NotIn", Values: []string{"v3"}},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := convertNodeSelectorRequirements(tt.input)
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("convertNodeSelectorRequirements() = %#v, want %#v", got, tt.want)
}
})
}
}

0 comments on commit 8948b24

Please sign in to comment.