Skip to content

Commit

Permalink
test(propeller): Update unittests
Browse files Browse the repository at this point in the history 
Update unittests corresponding to the changes from SharedIndexInformer to Informer

Resolves: flyteorg#5087
Signed-off-by: Chi-Sheng Liu <[email protected]>
  • Loading branch information
@MortalHappiness
MortalHappiness committed Apr 4, 2024
1 parent c37f0a4 commit 70d0573
Show file tree
Hide file tree
Showing 5 changed files with 600 additions and 51 deletions.
25 changes: 15 additions & 10 deletions flytepropeller/pkg/controller/nodes/task/k8s/plugin_collector.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,20 @@ package k8s

import (
"context"
"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
corev1 "k8s.io/api/core/v1"
"runtime/pprof"
"strings"
"sync"
"time"

corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/runtime/schema"
"sigs.k8s.io/controller-runtime/pkg/cache"

"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
"github.com/flyteorg/flyte/flytestdlib/contextutils"
"github.com/flyteorg/flyte/flytestdlib/logger"
"github.com/flyteorg/flyte/flytestdlib/promutils"
"github.com/flyteorg/flyte/flytestdlib/promutils/labeled"
"k8s.io/apimachinery/pkg/runtime/schema"
"sigs.k8s.io/controller-runtime/pkg/cache"
)

const resourceLevelMonitorCycleDuration = 10 * time.Second
Expand Down Expand Up @@ -52,7 +53,7 @@ type ResourceLevelMonitor struct {
// K8s resource created by a plugin will have (as yet, Flyte doesn't have any plugins that create cluster level resources and
// it probably won't for a long time). We can't assume that all the operators and CRDs that Flyte will ever work with will have
// the exact same set of labels or annotations or owner references. The only thing we can really count on is namespace.
func (r *ResourceLevelMonitor) countList(ctx context.Context, pods *corev1.PodList) map[string]int {
func (r *ResourceLevelMonitor) countList(pods *corev1.PodList) map[string]int {
// Map of namespace to counts
counts := map[string]int{}

Expand All @@ -64,6 +65,13 @@ func (r *ResourceLevelMonitor) countList(ctx context.Context, pods *corev1.PodLi
return counts
}

func (r *ResourceLevelMonitor) setLevels(ctx context.Context, counts map[string]int) {
for ns, count := range counts {
withNamespaceCtx := contextutils.WithNamespace(ctx, ns)
r.Levels.Set(withNamespaceCtx, float64(count))
}
}

// The context here is expected to already have a value for the KindKey
func (r *ResourceLevelMonitor) collect(ctx context.Context) {
// Emit gauges at the namespace layer - since these are just K8s resources, we cannot be guaranteed to have the necessary
Expand All @@ -75,12 +83,9 @@ func (r *ResourceLevelMonitor) collect(ctx context.Context) {
return
}

counts := r.countList(ctx, pods)
counts := r.countList(pods)

for ns, count := range counts {
withNamespaceCtx := contextutils.WithNamespace(ctx, ns)
r.Levels.Set(withNamespaceCtx, float64(count))
}
r.setLevels(ctx, counts)
}

func (r *ResourceLevelMonitor) RunCollector(ctx context.Context) {
Expand Down
62 changes: 34 additions & 28 deletions flytepropeller/pkg/controller/nodes/task/k8s/plugin_collector_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,40 +8,44 @@ import (

"github.com/prometheus/client_golang/prometheus/testutil"
"github.com/stretchr/testify/assert"
v1 "k8s.io/api/core/v1"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes/scheme"
"k8s.io/client-go/tools/cache"

"github.com/flyteorg/flyte/flytepropeller/pkg/controller/executors/mocks"
"github.com/flyteorg/flyte/flytestdlib/promutils"
)

var pods = []interface{}{
&v1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "a",
Namespace: "ns-a",
// Set variable podList using v1.PodList
var pods = &corev1.PodList{
Items: []corev1.Pod{
{
ObjectMeta: metav1.ObjectMeta{
Name: "a",
Namespace: "ns-a",
},
},
},
&v1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "b",
Namespace: "ns-a",
{
ObjectMeta: metav1.ObjectMeta{
Name: "b",
Namespace: "ns-a",
},
},
},
&v1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "c",
Namespace: "ns-b",
{
ObjectMeta: metav1.ObjectMeta{
Name: "c",
Namespace: "ns-b",
},
},
},
}

func TestNewResourceLevelMonitor(t *testing.T) {
x := v1.Pod{}
x := corev1.Pod{}
x.GetObjectMeta()
lm := ResourceLevelMonitor{}
res := lm.countList(context.Background(), pods)
res := lm.countList(pods)
assert.Equal(t, 2, res["ns-a"])
assert.Equal(t, 1, res["ns-b"])
}
Expand All @@ -63,23 +67,24 @@ type MyFakeStore struct {
cache.Store
}

func (m MyFakeStore) List() []interface{} {
return pods
}

func TestResourceLevelMonitor_collect(t *testing.T) {
func TestResourceLevelMonitor_setLevels(t *testing.T) {
ctx := context.Background()
scope := promutils.NewScope("testscope")
fakeKubeClient := mocks.NewFakeKubeClient()

kinds, _, err := scheme.Scheme.ObjectKinds(&v1.Pod{})
kinds, _, err := scheme.Scheme.ObjectKinds(&corev1.Pod{})
assert.NoError(t, err)
myInformer := MyFakeInformer{
store: MyFakeStore{},
}

index := NewResourceMonitorIndex()
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, myInformer, kinds[0])
rm.collect(ctx)
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, myInformer, kinds[0], fakeKubeClient)
counts := map[string]int{
"ns-a": 2,
"ns-b": 1,
}
rm.setLevels(ctx, counts)

var expected = `
# HELP testscope:k8s_resources Current levels of K8s objects as seen from their informer caches
Expand All @@ -95,15 +100,16 @@ func TestResourceLevelMonitor_collect(t *testing.T) {
func TestResourceLevelMonitorSingletonness(t *testing.T) {
ctx := context.Background()
scope := promutils.NewScope("testscope")
fakeKubeClient := mocks.NewFakeKubeClient()

kinds, _, err := scheme.Scheme.ObjectKinds(&v1.Pod{})
kinds, _, err := scheme.Scheme.ObjectKinds(&corev1.Pod{})
assert.NoError(t, err)
myInformer := MyFakeInformer{
store: MyFakeStore{},
}

index := NewResourceMonitorIndex()
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, myInformer, kinds[0])
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, myInformer, kinds[0], fakeKubeClient)
fmt.Println(rm)
//rm2 := index.GetOrCreateResourceLevelMonitor(ctx, scope, myInformer, kinds[0])

Expand Down
2 changes: 1 addition & 1 deletion flytepropeller/pkg/controller/nodes/task/k8s/plugin_manager_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -994,7 +994,7 @@ func TestResourceManagerConstruction(t *testing.T) {
si, err := getPluginInformer(ctx, fakeKubeClient, &v1.Pod{})
assert.NotNil(t, si)
assert.NoError(t, err)
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, si, gvk)
rm := index.GetOrCreateResourceLevelMonitor(ctx, scope, si, gvk, fakeKubeClient)
assert.NotNil(t, rm)
}

Expand Down
17 changes: 9 additions & 8 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ require (
github.com/cespare/xxhash/v2 v2.2.0 // indirect
github.com/cloudevents/sdk-go/binding/format/protobuf/v2 v2.14.0 // indirect
github.com/cloudevents/sdk-go/protocol/kafka_sarama/v2 v2.8.0 // indirect
github.com/cloudevents/sdk-go/v2 v2.14.0 // indirect
github.com/cloudevents/sdk-go/v2 v2.15.2 // indirect
github.com/containerd/containerd v1.5.10 // indirect
github.com/coocood/freecache v1.1.1 // indirect
github.com/coreos/go-oidc/v3 v3.6.0 // indirect
Expand All @@ -74,7 +74,7 @@ require (
github.com/eapache/go-resiliency v1.2.0 // indirect
github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21 // indirect
github.com/eapache/queue v1.1.0 // indirect
github.com/emicklei/go-restful/v3 v3.11.0 // indirect
github.com/emicklei/go-restful/v3 v3.12.0 // indirect
github.com/enescakir/emoji v1.0.0 // indirect
github.com/evanphx/json-patch v5.6.0+incompatible // indirect
github.com/evanphx/json-patch/v5 v5.9.0 // indirect
Expand All @@ -95,7 +95,7 @@ require (
github.com/go-ozzo/ozzo-validation/v4 v4.3.0 // indirect
github.com/go-redis/redis v6.15.7+incompatible // indirect
github.com/go-test/deep v1.0.7 // indirect
github.com/goccy/go-json v0.4.8 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/godbus/dbus/v5 v5.0.4 // indirect
github.com/gofrs/uuid v4.2.0+incompatible // indirect
github.com/gogo/protobuf v1.3.2 // indirect
Expand Down Expand Up @@ -149,11 +149,12 @@ require (
github.com/kubeflow/training-operator v1.5.0-rc.0 // indirect
github.com/kylelemons/godebug v1.1.0 // indirect
github.com/landoop/tableprinter v0.0.0-20180806200924-8bd8c2576d27 // indirect
github.com/lestrrat-go/backoff/v2 v2.0.7 // indirect
github.com/lestrrat-go/httpcc v1.0.0 // indirect
github.com/lestrrat-go/iter v1.0.1 // indirect
github.com/lestrrat-go/jwx v1.1.6 // indirect
github.com/lestrrat-go/option v1.0.0 // indirect
github.com/lestrrat-go/backoff/v2 v2.0.8 // indirect
github.com/lestrrat-go/blackmagic v1.0.2 // indirect
github.com/lestrrat-go/httpcc v1.0.1 // indirect
github.com/lestrrat-go/iter v1.0.2 // indirect
github.com/lestrrat-go/jwx v1.2.29 // indirect
github.com/lestrrat-go/option v1.0.1 // indirect
github.com/magiconair/properties v1.8.6 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/mattn/go-colorable v0.1.12 // indirect
Expand Down
Loading

0 comments on commit 70d0573

Please sign in to comment.