Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion controllers/state_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -438,7 +438,8 @@ func (n *ClusterPolicyController) applyDriverAutoUpgradeAnnotation() error {
updateRequired := false
value := "true"
annotationValue, annotationExists := node.Annotations[driverAutoUpgradeAnnotationKey]
if n.singleton.Spec.Driver.UpgradePolicy != nil &&
if n.singleton.Spec.Driver.IsEnabled() &&
n.singleton.Spec.Driver.UpgradePolicy != nil &&
n.singleton.Spec.Driver.UpgradePolicy.AutoUpgrade &&
!n.sandboxEnabled {
// check if we need to add the annotation
Expand Down
68 changes: 68 additions & 0 deletions controllers/state_manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ package controllers
import (
"testing"

upgrade_v1alpha1 "github.com/NVIDIA/k8s-operator-libs/api/upgrade/v1alpha1"
corev1 "k8s.io/api/core/v1"

gpuv1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1"
Expand Down Expand Up @@ -186,3 +187,70 @@ func TestHasMIGCapableGPU(t *testing.T) {
}
}
}

func TestShouldApplyDriverAutoUpgradeAnnotation(t *testing.T) {
tests := []struct {
name string
driverEnabled bool
autoUpgradeEnabled bool
sandboxEnabled bool
wantAnnotation bool
}{
{
name: "driver enabled with auto-upgrade",
driverEnabled: true,
autoUpgradeEnabled: true,
sandboxEnabled: false,
wantAnnotation: true,
},
{
name: "driver disabled with auto-upgrade enabled - should NOT apply annotation",
driverEnabled: false,
autoUpgradeEnabled: true,
sandboxEnabled: false,
wantAnnotation: false,
},
{
name: "driver enabled but auto-upgrade disabled",
driverEnabled: true,
autoUpgradeEnabled: false,
sandboxEnabled: false,
wantAnnotation: false,
},
{
name: "driver enabled with auto-upgrade but sandbox enabled",
driverEnabled: true,
autoUpgradeEnabled: true,
sandboxEnabled: true,
wantAnnotation: false,
},
{
name: "all disabled",
driverEnabled: false,
autoUpgradeEnabled: false,
sandboxEnabled: false,
wantAnnotation: false,
},
}

for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
// Create a mock ClusterPolicy
clusterPolicy := &gpuv1.ClusterPolicy{}
clusterPolicy.Spec.Driver.Enabled = &tc.driverEnabled
clusterPolicy.Spec.Driver.UpgradePolicy = &upgrade_v1alpha1.DriverUpgradePolicySpec{
AutoUpgrade: tc.autoUpgradeEnabled,
}

// Simulate the logic from applyDriverAutoUpgradeAnnotation
shouldApply := clusterPolicy.Spec.Driver.IsEnabled() &&
clusterPolicy.Spec.Driver.UpgradePolicy != nil &&
clusterPolicy.Spec.Driver.UpgradePolicy.AutoUpgrade &&
!tc.sandboxEnabled

if shouldApply != tc.wantAnnotation {
t.Errorf("Expected annotation to be applied: %v, but got: %v", tc.wantAnnotation, shouldApply)
}
})
}
}