mirror of
https://github.com/netbirdio/netbird.git
synced 2024-11-07 08:44:07 +01:00
chore: pkg import only once (#1222)
Signed-off-by: guoguangwu <guoguangwu@magic-shield.com>
This commit is contained in:
parent
46f5f148da
commit
73e57f17ea
@ -16,7 +16,6 @@ import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/encryption"
|
||||
"github.com/netbirdio/netbird/management/proto"
|
||||
mgmtProto "github.com/netbirdio/netbird/management/proto"
|
||||
mgmt "github.com/netbirdio/netbird/management/server"
|
||||
"github.com/netbirdio/netbird/management/server/mock_server"
|
||||
@ -95,8 +94,8 @@ func startMockManagement(t *testing.T) (*grpc.Server, net.Listener, *mock_server
|
||||
}
|
||||
|
||||
mgmtMockServer := &mock_server.ManagementServiceServerMock{
|
||||
GetServerKeyFunc: func(context.Context, *proto.Empty) (*proto.ServerKeyResponse, error) {
|
||||
response := &proto.ServerKeyResponse{
|
||||
GetServerKeyFunc: func(context.Context, *mgmtProto.Empty) (*mgmtProto.ServerKeyResponse, error) {
|
||||
response := &mgmtProto.ServerKeyResponse{
|
||||
Key: serverKey.PublicKey().String(),
|
||||
}
|
||||
return response, nil
|
||||
@ -300,19 +299,19 @@ func Test_SystemMetaDataFromClient(t *testing.T) {
|
||||
log.Fatalf("error while getting server public key from testclient, %v", err)
|
||||
}
|
||||
|
||||
var actualMeta *proto.PeerSystemMeta
|
||||
var actualMeta *mgmtProto.PeerSystemMeta
|
||||
var actualValidKey string
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
|
||||
mgmtMockServer.LoginFunc = func(ctx context.Context, msg *proto.EncryptedMessage) (*proto.EncryptedMessage, error) {
|
||||
mgmtMockServer.LoginFunc = func(ctx context.Context, msg *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
|
||||
peerKey, err := wgtypes.ParseKey(msg.GetWgPubKey())
|
||||
if err != nil {
|
||||
log.Warnf("error while parsing peer's Wireguard public key %s on Sync request.", msg.WgPubKey)
|
||||
return nil, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", msg.WgPubKey)
|
||||
}
|
||||
|
||||
loginReq := &proto.LoginRequest{}
|
||||
loginReq := &mgmtProto.LoginRequest{}
|
||||
err = encryption.DecryptMessage(peerKey, serverKey, msg.Body, loginReq)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
@ -322,7 +321,7 @@ func Test_SystemMetaDataFromClient(t *testing.T) {
|
||||
actualValidKey = loginReq.GetSetupKey()
|
||||
wg.Done()
|
||||
|
||||
loginResp := &proto.LoginResponse{}
|
||||
loginResp := &mgmtProto.LoginResponse{}
|
||||
encryptedResp, err := encryption.EncryptMessage(peerKey, serverKey, loginResp)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -343,7 +342,7 @@ func Test_SystemMetaDataFromClient(t *testing.T) {
|
||||
|
||||
wg.Wait()
|
||||
|
||||
expectedMeta := &proto.PeerSystemMeta{
|
||||
expectedMeta := &mgmtProto.PeerSystemMeta{
|
||||
Hostname: info.Hostname,
|
||||
GoOS: info.GoOS,
|
||||
Kernel: info.Kernel,
|
||||
@ -374,12 +373,12 @@ func Test_GetDeviceAuthorizationFlow(t *testing.T) {
|
||||
log.Fatalf("error while creating testClient: %v", err)
|
||||
}
|
||||
|
||||
expectedFlowInfo := &proto.DeviceAuthorizationFlow{
|
||||
expectedFlowInfo := &mgmtProto.DeviceAuthorizationFlow{
|
||||
Provider: 0,
|
||||
ProviderConfig: &proto.ProviderConfig{ClientID: "client"},
|
||||
ProviderConfig: &mgmtProto.ProviderConfig{ClientID: "client"},
|
||||
}
|
||||
|
||||
mgmtMockServer.GetDeviceAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*proto.EncryptedMessage, error) {
|
||||
mgmtMockServer.GetDeviceAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
|
||||
encryptedResp, err := encryption.EncryptMessage(serverKey, client.key, expectedFlowInfo)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -418,14 +417,14 @@ func Test_GetPKCEAuthorizationFlow(t *testing.T) {
|
||||
log.Fatalf("error while creating testClient: %v", err)
|
||||
}
|
||||
|
||||
expectedFlowInfo := &proto.PKCEAuthorizationFlow{
|
||||
ProviderConfig: &proto.ProviderConfig{
|
||||
expectedFlowInfo := &mgmtProto.PKCEAuthorizationFlow{
|
||||
ProviderConfig: &mgmtProto.ProviderConfig{
|
||||
ClientID: "client",
|
||||
ClientSecret: "secret",
|
||||
},
|
||||
}
|
||||
|
||||
mgmtMockServer.GetPKCEAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*proto.EncryptedMessage, error) {
|
||||
mgmtMockServer.GetPKCEAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
|
||||
encryptedResp, err := encryption.EncryptMessage(serverKey, client.key, expectedFlowInfo)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
Loading…
Reference in New Issue
Block a user