github.com/kubewharf/katalyst-core@v0.5.3/pkg/agent/orm/metamanager/manager_test.go (about)

     1  /*
     2  Copyright 2022 The Katalyst Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package metamanager
    18  
    19  import (
    20  	"io/ioutil"
    21  	"os"
    22  	"testing"
    23  
    24  	"github.com/stretchr/testify/assert"
    25  	"github.com/stretchr/testify/require"
    26  	v1 "k8s.io/api/core/v1"
    27  	v12 "k8s.io/apimachinery/pkg/apis/meta/v1"
    28  	"k8s.io/apimachinery/pkg/runtime"
    29  	"k8s.io/apimachinery/pkg/util/sets"
    30  
    31  	katalyst_base "github.com/kubewharf/katalyst-core/cmd/base"
    32  	"github.com/kubewharf/katalyst-core/cmd/katalyst-agent/app/options"
    33  	"github.com/kubewharf/katalyst-core/pkg/config"
    34  	"github.com/kubewharf/katalyst-core/pkg/metaserver"
    35  	"github.com/kubewharf/katalyst-core/pkg/metaserver/agent/pod"
    36  	"github.com/kubewharf/katalyst-core/pkg/metrics"
    37  )
    38  
    39  func generateTestMetaServer(conf *config.Configuration) (*metaserver.MetaServer, error) {
    40  	genericCtx, err := katalyst_base.GenerateFakeGenericContext([]runtime.Object{})
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  
    45  	return metaserver.NewMetaServer(genericCtx.Client, metrics.DummyMetrics{}, conf)
    46  }
    47  
    48  func generateTestConfiguration(checkpointDir string) *config.Configuration {
    49  	conf, _ := options.NewOptions().Config()
    50  
    51  	conf.MetaServerConfiguration.CheckpointManagerDir = checkpointDir
    52  
    53  	return conf
    54  }
    55  
    56  func TestReconcile(t *testing.T) {
    57  	t.Parallel()
    58  
    59  	ckDir, err := ioutil.TempDir("", "checkpoint-Test")
    60  	require.NoError(t, err)
    61  	defer func() { _ = os.RemoveAll(ckDir) }()
    62  
    63  	conf := generateTestConfiguration(ckDir)
    64  	metaServer, err := generateTestMetaServer(conf)
    65  	require.NoError(t, err)
    66  
    67  	metaServer.PodFetcher = &pod.PodFetcherStub{
    68  		PodList: []*v1.Pod{
    69  			{
    70  				ObjectMeta: v12.ObjectMeta{
    71  					Name: "pod0",
    72  					UID:  "pod0",
    73  				},
    74  			},
    75  			{
    76  				ObjectMeta: v12.ObjectMeta{
    77  					Name: "pod1",
    78  					UID:  "pod1",
    79  				},
    80  			},
    81  			{
    82  				ObjectMeta: v12.ObjectMeta{
    83  					Name: "pod2",
    84  					UID:  "pod2",
    85  				},
    86  			},
    87  		},
    88  	}
    89  
    90  	manager := NewManager(metrics.DummyMetrics{}, func() sets.String {
    91  		return sets.NewString("pod0", "pod3", "pod4", "pod5")
    92  	}, metaServer)
    93  
    94  	newPodList := make([]string, 0)
    95  	removePodList := make([]string, 0)
    96  
    97  	manager.RegistPodAddedFunc(func(podUID string) {
    98  		newPodList = append(newPodList, podUID)
    99  	})
   100  	manager.RegistPodDeletedFunc(func(podUID string) {
   101  		removePodList = append(removePodList, podUID)
   102  	})
   103  
   104  	manager.reconcile()
   105  	require.Equal(t, 2, len(newPodList))
   106  	require.Equal(t, 3, len(removePodList))
   107  }
   108  
   109  func TestReconcilePods(t *testing.T) {
   110  	t.Parallel()
   111  
   112  	ckDir, err := ioutil.TempDir("", "checkpoint-Test")
   113  	require.NoError(t, err)
   114  	defer func() { _ = os.RemoveAll(ckDir) }()
   115  
   116  	conf := generateTestConfiguration(ckDir)
   117  	metaServer, err := generateTestMetaServer(conf)
   118  	require.NoError(t, err)
   119  
   120  	metaServer.PodFetcher = &pod.PodFetcherStub{
   121  		PodList: []*v1.Pod{
   122  			{
   123  				ObjectMeta: v12.ObjectMeta{
   124  					Name: "pod0",
   125  					UID:  "pod0",
   126  				},
   127  			},
   128  			{
   129  				ObjectMeta: v12.ObjectMeta{
   130  					Name: "pod1",
   131  					UID:  "pod1",
   132  				},
   133  			},
   134  			{
   135  				ObjectMeta: v12.ObjectMeta{
   136  					Name: "pod2",
   137  					UID:  "pod2",
   138  				},
   139  			},
   140  		},
   141  	}
   142  
   143  	manager := NewManager(metrics.DummyMetrics{}, func() sets.String {
   144  		return sets.NewString("pod0", "pod3", "pod4", "pod5")
   145  	}, metaServer)
   146  
   147  	p := manager.GetPods()
   148  	assert.Equal(t, 3, len(p))
   149  
   150  	newPods, removePods, err := manager.ReconcilePods()
   151  	assert.NoError(t, err)
   152  	assert.Equal(t, 2, len(newPods))
   153  	assert.Equal(t, 3, len(removePods))
   154  }