mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-02 20:09:57 +08:00
84 lines
2.5 KiB
Go
84 lines
2.5 KiB
Go
/*
|
|
* Licensed to the LF AI & Data foundation under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you 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 contextutil
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"google.golang.org/grpc/metadata"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util"
|
|
"github.com/milvus-io/milvus/pkg/util/crypto"
|
|
)
|
|
|
|
func TestAppendToIncomingContext(t *testing.T) {
|
|
t.Run("invalid kvs", func(t *testing.T) {
|
|
assert.Panics(t, func() {
|
|
// nolint
|
|
AppendToIncomingContext(context.Background(), "foo")
|
|
})
|
|
})
|
|
|
|
t.Run("valid kvs", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctx = AppendToIncomingContext(ctx, "foo", "bar")
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "bar", md.Get("foo")[0])
|
|
})
|
|
}
|
|
|
|
func TestGetCurUserFromContext(t *testing.T) {
|
|
_, err := GetCurUserFromContext(context.Background())
|
|
assert.Error(t, err)
|
|
|
|
_, err = GetCurUserFromContext(metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{})))
|
|
assert.Error(t, err)
|
|
|
|
_, err = GetCurUserFromContext(GetContext(context.Background(), "123456"))
|
|
assert.Error(t, err)
|
|
|
|
root := "root"
|
|
password := "123456"
|
|
username, err := GetCurUserFromContext(GetContext(context.Background(), fmt.Sprintf("%s%s%s", root, util.CredentialSeperator, password)))
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, root, username)
|
|
|
|
{
|
|
u, p, e := GetAuthInfoFromContext(GetContext(context.Background(), fmt.Sprintf("%s%s%s", root, util.CredentialSeperator, password)))
|
|
assert.NoError(t, e)
|
|
assert.Equal(t, "root", u)
|
|
assert.Equal(t, password, p)
|
|
}
|
|
}
|
|
|
|
func GetContext(ctx context.Context, originValue string) context.Context {
|
|
authKey := strings.ToLower(util.HeaderAuthorize)
|
|
authValue := crypto.Base64Encode(originValue)
|
|
contextMap := map[string]string{
|
|
authKey: authValue,
|
|
}
|
|
md := metadata.New(contextMap)
|
|
return metadata.NewIncomingContext(ctx, md)
|
|
}
|