mirror of
https://github.com/fankes/beszel.git
synced 2025-10-20 02:09:28 +08:00
rename /src to /internal (sorry i'll fix the prs)
This commit is contained in:
henrygd
336
internal/cmd/agent/agent_test.go
Normal file
336
internal/cmd/agent/agent_test.go
Normal file
@@ -0,0 +1,336 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/henrygd/beszel/agent"
|
||||
|
||||
"github.com/spf13/pflag"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
func TestGetAddress(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "default port when no config",
|
||||
opts: cmdOptions{},
|
||||
expected: ":45876",
|
||||
},
|
||||
{
|
||||
name: "use address from flag",
|
||||
opts: cmdOptions{
|
||||
listen: "8080",
|
||||
},
|
||||
expected: ":8080",
|
||||
},
|
||||
{
|
||||
name: "use unix socket from flag",
|
||||
opts: cmdOptions{
|
||||
listen: "/tmp/beszel.sock",
|
||||
},
|
||||
expected: "/tmp/beszel.sock",
|
||||
},
|
||||
{
|
||||
name: "use LISTEN env var",
|
||||
opts: cmdOptions{},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": "1.2.3.4:9090",
|
||||
},
|
||||
expected: "1.2.3.4:9090",
|
||||
},
|
||||
{
|
||||
name: "use legacy PORT env var",
|
||||
opts: cmdOptions{},
|
||||
envVars: map[string]string{
|
||||
"PORT": "7070",
|
||||
},
|
||||
expected: ":7070",
|
||||
},
|
||||
{
|
||||
name: "use unix socket from env var",
|
||||
opts: cmdOptions{
|
||||
listen: "",
|
||||
},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": "/tmp/beszel.sock",
|
||||
},
|
||||
expected: "/tmp/beszel.sock",
|
||||
},
|
||||
{
|
||||
name: "flag takes precedence over env vars",
|
||||
opts: cmdOptions{
|
||||
listen: ":8080",
|
||||
},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": ":9090",
|
||||
"PORT": "7070",
|
||||
},
|
||||
expected: ":8080",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Setup environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
|
||||
addr := tt.opts.getAddress()
|
||||
assert.Equal(t, tt.expected, addr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadPublicKeys(t *testing.T) {
|
||||
// Generate a test key
|
||||
_, priv, err := ed25519.GenerateKey(nil)
|
||||
require.NoError(t, err)
|
||||
signer, err := ssh.NewSignerFromKey(priv)
|
||||
require.NoError(t, err)
|
||||
pubKey := ssh.MarshalAuthorizedKey(signer.PublicKey())
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
setupFiles map[string][]byte
|
||||
wantErr bool
|
||||
errContains string
|
||||
}{
|
||||
{
|
||||
name: "load key from flag",
|
||||
opts: cmdOptions{
|
||||
key: string(pubKey),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "load key from env var",
|
||||
envVars: map[string]string{
|
||||
"KEY": string(pubKey),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "load key from file",
|
||||
envVars: map[string]string{
|
||||
"KEY_FILE": "testkey.pub",
|
||||
},
|
||||
setupFiles: map[string][]byte{
|
||||
"testkey.pub": pubKey,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "error when no key provided",
|
||||
wantErr: true,
|
||||
errContains: "no key provided",
|
||||
},
|
||||
{
|
||||
name: "error on invalid key file",
|
||||
envVars: map[string]string{
|
||||
"KEY_FILE": "nonexistent.pub",
|
||||
},
|
||||
wantErr: true,
|
||||
errContains: "failed to read key file",
|
||||
},
|
||||
{
|
||||
name: "error on invalid key data",
|
||||
opts: cmdOptions{
|
||||
key: "invalid-key-data",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create a temporary directory for test files
|
||||
if len(tt.setupFiles) > 0 {
|
||||
tmpDir := t.TempDir()
|
||||
for name, content := range tt.setupFiles {
|
||||
path := filepath.Join(tmpDir, name)
|
||||
err := os.WriteFile(path, content, 0600)
|
||||
require.NoError(t, err)
|
||||
if tt.envVars != nil {
|
||||
tt.envVars["KEY_FILE"] = path
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set up environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
|
||||
keys, err := tt.opts.loadPublicKeys()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
if tt.errContains != "" {
|
||||
assert.Contains(t, err.Error(), tt.errContains)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, keys, 1)
|
||||
assert.Equal(t, signer.PublicKey().Type(), keys[0].Type())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetNetwork(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "NETWORK env var",
|
||||
envVars: map[string]string{
|
||||
"NETWORK": "tcp4",
|
||||
},
|
||||
expected: "tcp4",
|
||||
},
|
||||
{
|
||||
name: "only port",
|
||||
opts: cmdOptions{listen: "8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "ipv4 address",
|
||||
opts: cmdOptions{listen: "1.2.3.4:8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "ipv6 address",
|
||||
opts: cmdOptions{listen: "[2001:db8::1]:8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "unix network",
|
||||
opts: cmdOptions{listen: "/tmp/beszel.sock"},
|
||||
expected: "unix",
|
||||
},
|
||||
{
|
||||
name: "env var network",
|
||||
opts: cmdOptions{listen: ":8080"},
|
||||
envVars: map[string]string{"NETWORK": "tcp4"},
|
||||
expected: "tcp4",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Setup environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
network := agent.GetNetwork(tt.opts.listen)
|
||||
assert.Equal(t, tt.expected, network)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseFlags(t *testing.T) {
|
||||
// Save original command line arguments and restore after test
|
||||
oldArgs := os.Args
|
||||
defer func() {
|
||||
os.Args = oldArgs
|
||||
pflag.CommandLine = pflag.NewFlagSet(os.Args[0], pflag.ExitOnError)
|
||||
}()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
args []string
|
||||
expected cmdOptions
|
||||
}{
|
||||
{
|
||||
name: "no flags",
|
||||
args: []string{"cmd"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag only",
|
||||
args: []string{"cmd", "-key", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag double dash",
|
||||
args: []string{"cmd", "--key", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag short",
|
||||
args: []string{"cmd", "-k", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag only",
|
||||
args: []string{"cmd", "-listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag double dash",
|
||||
args: []string{"cmd", "--listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag short",
|
||||
args: []string{"cmd", "-l", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "both flags",
|
||||
args: []string{"cmd", "-key", "testkey", "-listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Reset flags for each test
|
||||
pflag.CommandLine = pflag.NewFlagSet(tt.args[0], pflag.ExitOnError)
|
||||
os.Args = tt.args
|
||||
|
||||
var opts cmdOptions
|
||||
opts.parse()
|
||||
pflag.Parse()
|
||||
|
||||
assert.Equal(t, tt.expected, opts)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user