mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-04 21:09:06 +08:00
95 lines
2.7 KiB
Go
95 lines
2.7 KiB
Go
|
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||
|
|
//
|
||
|
|
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||
|
|
// with the License. You may obtain a copy of the License at
|
||
|
|
//
|
||
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
//
|
||
|
|
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||
|
|
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||
|
|
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||
|
|
|
||
|
|
package masterservice
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"path"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/milvus-io/milvus/internal/util/sessionutil"
|
||
|
|
"github.com/milvus-io/milvus/internal/util/typeutil"
|
||
|
|
"github.com/stretchr/testify/assert"
|
||
|
|
"go.etcd.io/etcd/clientv3"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestProxyNodeManager(t *testing.T) {
|
||
|
|
Params.Init()
|
||
|
|
cli, err := clientv3.New(clientv3.Config{Endpoints: []string{Params.EtcdAddress}})
|
||
|
|
assert.Nil(t, err)
|
||
|
|
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
sessKey := path.Join(Params.MetaRootPath, sessionutil.DefaultServiceRoot)
|
||
|
|
cli.Delete(ctx, sessKey, clientv3.WithPrefix())
|
||
|
|
defer cli.Delete(ctx, sessKey, clientv3.WithPrefix())
|
||
|
|
s1 := sessionutil.Session{
|
||
|
|
ServerID: 100,
|
||
|
|
}
|
||
|
|
b1, err := json.Marshal(&s1)
|
||
|
|
assert.Nil(t, err)
|
||
|
|
k1 := path.Join(sessKey, typeutil.ProxyNodeRole+"-100")
|
||
|
|
_, err = cli.Put(ctx, k1, string(b1))
|
||
|
|
assert.Nil(t, err)
|
||
|
|
|
||
|
|
s0 := sessionutil.Session{
|
||
|
|
ServerID: 99,
|
||
|
|
}
|
||
|
|
b0, err := json.Marshal(&s0)
|
||
|
|
assert.Nil(t, err)
|
||
|
|
k0 := path.Join(sessKey, typeutil.ProxyNodeRole+"-99")
|
||
|
|
_, err = cli.Put(ctx, k0, string(b0))
|
||
|
|
assert.Nil(t, err)
|
||
|
|
|
||
|
|
f1 := func(sess []*sessionutil.Session) {
|
||
|
|
assert.Equal(t, len(sess), 2)
|
||
|
|
assert.Equal(t, int64(100), sess[0].ServerID)
|
||
|
|
assert.Equal(t, int64(99), sess[1].ServerID)
|
||
|
|
t.Log("get sessions", sess[0], sess[1])
|
||
|
|
}
|
||
|
|
|
||
|
|
pm, err := newProxyNodeManager(ctx, []string{Params.EtcdAddress}, f1)
|
||
|
|
assert.Nil(t, err)
|
||
|
|
fa := func(sess *sessionutil.Session) {
|
||
|
|
assert.Equal(t, int64(101), sess.ServerID)
|
||
|
|
t.Log("add session", sess)
|
||
|
|
}
|
||
|
|
fd := func(sess *sessionutil.Session) {
|
||
|
|
assert.Equal(t, int64(100), sess.ServerID)
|
||
|
|
t.Log("del session", sess)
|
||
|
|
}
|
||
|
|
pm.AddSession(fa)
|
||
|
|
pm.DelSession(fd)
|
||
|
|
|
||
|
|
err = pm.WatchProxyNode()
|
||
|
|
assert.Nil(t, err)
|
||
|
|
t.Log("======== start watch proxy node ==========")
|
||
|
|
|
||
|
|
s2 := sessionutil.Session{
|
||
|
|
ServerID: 101,
|
||
|
|
}
|
||
|
|
b2, err := json.Marshal(&s2)
|
||
|
|
assert.Nil(t, err)
|
||
|
|
k2 := path.Join(sessKey, typeutil.ProxyNodeRole+"-101")
|
||
|
|
_, err = cli.Put(ctx, k2, string(b2))
|
||
|
|
assert.Nil(t, err)
|
||
|
|
|
||
|
|
_, err = cli.Delete(ctx, k1)
|
||
|
|
assert.Nil(t, err)
|
||
|
|
time.Sleep(time.Second)
|
||
|
|
pm.Stop()
|
||
|
|
time.Sleep(time.Second)
|
||
|
|
}
|