2019-03-13 04:14:30 -07:00
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
2021-02-25 00:01:08 -08:00
|
|
|
"errors"
|
2019-03-13 04:14:30 -07:00
|
|
|
"testing"
|
|
|
|
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
)
|
|
|
|
|
2021-02-25 00:01:08 -08:00
|
|
|
func TestPropertySessionSegment(t *testing.T) {
|
|
|
|
cases := []struct {
|
|
|
|
Case string
|
|
|
|
ExpectedEnabled bool
|
|
|
|
ExpectedString string
|
|
|
|
UserName string
|
|
|
|
Host string
|
|
|
|
DefaultUserName string
|
|
|
|
DefaultUserNameEnv string
|
|
|
|
SSHSession bool
|
|
|
|
SSHClient bool
|
|
|
|
Root bool
|
|
|
|
DisplayUser bool
|
|
|
|
DisplayHost bool
|
|
|
|
DisplayDefault bool
|
|
|
|
HostColor string
|
|
|
|
UserColor string
|
|
|
|
GOOS string
|
|
|
|
HostError bool
|
|
|
|
}{
|
|
|
|
{
|
|
|
|
Case: "user and computer",
|
|
|
|
ExpectedString: "john at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "user and computer with host color",
|
|
|
|
ExpectedString: "john at <yellow>company-laptop</>",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
HostColor: "yellow",
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "user and computer with user color",
|
|
|
|
ExpectedString: "<yellow>john</> at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
UserColor: "yellow",
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "user and computer with both colors",
|
|
|
|
ExpectedString: "<yellow>john</> at <green>company-laptop</>",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
UserColor: "yellow",
|
|
|
|
HostColor: "green",
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "SSH Session",
|
|
|
|
ExpectedString: "ssh john at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
SSHSession: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "SSH Client",
|
|
|
|
ExpectedString: "ssh john at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
SSHClient: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "SSH Client",
|
|
|
|
ExpectedString: "ssh john at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
DisplayUser: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
UserName: "john",
|
|
|
|
SSHClient: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "only user name",
|
|
|
|
ExpectedString: "john",
|
|
|
|
Host: "company-laptop",
|
|
|
|
UserName: "john",
|
|
|
|
DisplayUser: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "windows user name",
|
|
|
|
ExpectedString: "john at company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
UserName: "surface\\john",
|
|
|
|
DisplayHost: true,
|
|
|
|
DisplayUser: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
GOOS: string(Windows),
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "only host name",
|
|
|
|
ExpectedString: "company-laptop",
|
|
|
|
Host: "company-laptop",
|
|
|
|
UserName: "john",
|
|
|
|
DisplayDefault: true,
|
|
|
|
DisplayHost: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "display default - hidden",
|
|
|
|
Host: "company-laptop",
|
|
|
|
UserName: "john",
|
|
|
|
DefaultUserName: "john",
|
|
|
|
DisplayDefault: false,
|
|
|
|
DisplayHost: true,
|
|
|
|
DisplayUser: true,
|
|
|
|
ExpectedEnabled: false,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "display default with env var - hidden",
|
|
|
|
Host: "company-laptop",
|
|
|
|
UserName: "john",
|
|
|
|
DefaultUserNameEnv: "john",
|
|
|
|
DefaultUserName: "jake",
|
|
|
|
DisplayDefault: false,
|
|
|
|
DisplayHost: true,
|
|
|
|
DisplayUser: true,
|
|
|
|
ExpectedEnabled: false,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "host error",
|
|
|
|
ExpectedString: "john at unknown",
|
|
|
|
Host: "company-laptop",
|
|
|
|
HostError: true,
|
|
|
|
UserName: "john",
|
|
|
|
DisplayHost: true,
|
|
|
|
DisplayUser: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
2020-11-03 11:06:11 -08:00
|
|
|
}
|
|
|
|
|
2021-02-25 00:01:08 -08:00
|
|
|
for _, tc := range cases {
|
|
|
|
env := new(MockedEnvironment)
|
|
|
|
env.On("getCurrentUser", nil).Return(tc.UserName)
|
|
|
|
env.On("getRuntimeGOOS", nil).Return(tc.GOOS)
|
|
|
|
if tc.HostError {
|
|
|
|
env.On("getHostName", nil).Return(tc.Host, errors.New("oh snap"))
|
|
|
|
} else {
|
|
|
|
env.On("getHostName", nil).Return(tc.Host, nil)
|
|
|
|
}
|
|
|
|
var SSHSession string
|
|
|
|
if tc.SSHSession {
|
|
|
|
SSHSession = "zezzion"
|
|
|
|
}
|
|
|
|
var SSHClient string
|
|
|
|
if tc.SSHClient {
|
|
|
|
SSHClient = "clientz"
|
|
|
|
}
|
|
|
|
env.On("getenv", "SSH_CONNECTION").Return(SSHSession)
|
|
|
|
env.On("getenv", "SSH_CLIENT").Return(SSHClient)
|
|
|
|
env.On("getenv", "SSH_CLIENT").Return(SSHSession)
|
|
|
|
env.On("getenv", defaultUserEnvVar).Return(tc.DefaultUserNameEnv)
|
|
|
|
env.On("isRunningAsRoot", nil).Return(tc.Root)
|
|
|
|
props := &properties{
|
|
|
|
values: map[Property]interface{}{
|
|
|
|
UserInfoSeparator: " at ",
|
|
|
|
SSHIcon: "ssh ",
|
|
|
|
DefaultUserName: tc.DefaultUserName,
|
|
|
|
DisplayDefault: tc.DisplayDefault,
|
|
|
|
DisplayUser: tc.DisplayUser,
|
|
|
|
DisplayHost: tc.DisplayHost,
|
|
|
|
HostColor: tc.HostColor,
|
|
|
|
UserColor: tc.UserColor,
|
|
|
|
},
|
|
|
|
}
|
|
|
|
session := &session{
|
|
|
|
env: env,
|
|
|
|
props: props,
|
|
|
|
}
|
|
|
|
assert.Equal(t, tc.ExpectedEnabled, session.enabled(), tc.Case)
|
|
|
|
if tc.ExpectedEnabled {
|
|
|
|
assert.Equal(t, tc.ExpectedString, session.string(), tc.Case)
|
|
|
|
}
|
2020-10-12 03:53:54 -07:00
|
|
|
}
|
|
|
|
}
|
2021-02-15 13:01:09 -08:00
|
|
|
|
|
|
|
func TestSessionSegmentTemplate(t *testing.T) {
|
|
|
|
cases := []struct {
|
2021-02-24 00:04:04 -08:00
|
|
|
Case string
|
|
|
|
ExpectedEnabled bool
|
|
|
|
ExpectedString string
|
|
|
|
UserName string
|
|
|
|
DefaultUserName string
|
|
|
|
ComputerName string
|
|
|
|
SSHSession bool
|
|
|
|
Root bool
|
|
|
|
Template string
|
2021-02-15 13:01:09 -08:00
|
|
|
}{
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "user and computer",
|
|
|
|
ExpectedString: "john@company-laptop",
|
|
|
|
ComputerName: "company-laptop",
|
|
|
|
UserName: "john",
|
|
|
|
Template: "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "user only",
|
|
|
|
ExpectedString: "john",
|
|
|
|
UserName: "john",
|
|
|
|
Template: "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "user with ssh",
|
|
|
|
ExpectedString: "john on remote",
|
|
|
|
UserName: "john",
|
|
|
|
SSHSession: true,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Template: "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "user without ssh",
|
|
|
|
ExpectedString: "john",
|
|
|
|
UserName: "john",
|
|
|
|
SSHSession: false,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Template: "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "user with root and ssh",
|
|
|
|
ExpectedString: "super john on remote",
|
|
|
|
UserName: "john",
|
|
|
|
SSHSession: true,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Root: true,
|
|
|
|
Template: "{{if .Root}}super {{end}}{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
{
|
2021-02-24 00:04:04 -08:00
|
|
|
Case: "no template",
|
2021-02-25 00:01:08 -08:00
|
|
|
ExpectedString: "\uf817 john@remote",
|
2021-02-24 00:04:04 -08:00
|
|
|
UserName: "john",
|
|
|
|
SSHSession: true,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Root: true,
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "default user not equal",
|
|
|
|
ExpectedString: "john",
|
|
|
|
UserName: "john",
|
|
|
|
DefaultUserName: "jack",
|
|
|
|
SSHSession: true,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Root: true,
|
|
|
|
Template: "{{if ne .DefaultUserName .UserName}}{{.UserName}}{{end}}",
|
|
|
|
ExpectedEnabled: true,
|
|
|
|
},
|
|
|
|
{
|
|
|
|
Case: "default user equal",
|
|
|
|
ExpectedString: "",
|
|
|
|
UserName: "john",
|
|
|
|
DefaultUserName: "john",
|
|
|
|
SSHSession: true,
|
|
|
|
ComputerName: "remote",
|
|
|
|
Root: true,
|
|
|
|
Template: "{{if ne .DefaultUserName .UserName}}{{.UserName}}{{end}}",
|
|
|
|
ExpectedEnabled: false,
|
2021-02-15 13:01:09 -08:00
|
|
|
},
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, tc := range cases {
|
|
|
|
env := new(MockedEnvironment)
|
|
|
|
env.On("getCurrentUser", nil).Return(tc.UserName)
|
|
|
|
env.On("getRuntimeGOOS", nil).Return("burp")
|
|
|
|
env.On("getHostName", nil).Return(tc.ComputerName, nil)
|
|
|
|
var SSHSession string
|
|
|
|
if tc.SSHSession {
|
|
|
|
SSHSession = "zezzion"
|
|
|
|
}
|
|
|
|
env.On("getenv", "SSH_CONNECTION").Return(SSHSession)
|
|
|
|
env.On("getenv", "SSH_CLIENT").Return(SSHSession)
|
|
|
|
env.On("isRunningAsRoot", nil).Return(tc.Root)
|
2021-02-24 00:04:04 -08:00
|
|
|
env.On("getenv", defaultUserEnvVar).Return(tc.DefaultUserName)
|
2021-02-15 13:01:09 -08:00
|
|
|
props := &properties{
|
|
|
|
values: map[Property]interface{}{
|
|
|
|
SegmentTemplate: tc.Template,
|
|
|
|
},
|
|
|
|
}
|
|
|
|
session := &session{
|
|
|
|
env: env,
|
|
|
|
props: props,
|
|
|
|
}
|
2021-02-24 00:04:04 -08:00
|
|
|
assert.Equal(t, tc.ExpectedEnabled, session.enabled(), tc.Case)
|
|
|
|
if tc.ExpectedEnabled {
|
|
|
|
assert.Equal(t, tc.ExpectedString, session.string(), tc.Case)
|
|
|
|
}
|
2021-02-15 13:01:09 -08:00
|
|
|
}
|
|
|
|
}
|