feat(session): default user override

resolves #447
JanDeDobbeleer · Feb 24, 2021 · 58f1789 · 58f1789
1 parent 1a907f3
commit 58f1789
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 60 deletions.
diff --git a/docs/docs/segment-session.md b/docs/docs/segment-session.md
@@ -39,8 +39,13 @@ properties below. Only used when a value is set, making the above properties obs
 ## Template Properties
 
 - `.UserName`: `string` - the current user's name
+- `.DefaultUserName`: - the default user name (set with the `POSH_SESSION_DEFAULT_USER` env var or `default_user_name` property)
 - `.ComputerName`: `string` - the current computer's name
 - `.SSHSession`: `boolean` - active SSH session or not
 - `.Root`: `boolean` - are you a root/admin user or not
 
+## Environmnent Variables
+
+- `POSH_SESSION_DEFAULT_USER` - used to override the hardcoded `default_user_name` property
+
 [colors]: /docs/configure#colors
diff --git a/src/segment_session.go b/src/segment_session.go
@@ -6,12 +6,14 @@ import (
 )
 
 type session struct {
-	props        *properties
-	env          environmentInfo
-	UserName     string
-	ComputerName string
-	SSHSession   bool
-	Root         bool
+	props           *properties
+	env             environmentInfo
+	UserName        string
+	DefaultUserName string
+	ComputerName    string
+	SSHSession      bool
+	Root            bool
+	templateText    string
 }
 
 const (
@@ -29,12 +31,26 @@ const (
 	SSHIcon Property = "ssh_icon"
 	// DefaultUserName holds the default user of the platform
 	DefaultUserName Property = "default_user_name"
+
+	defaultUserEnvVar = "POSH_SESSION_DEFAULT_USER"
 )
 
 func (s *session) enabled() bool {
 	s.UserName = s.getUserName()
+	s.ComputerName = s.getComputerName()
+	s.SSHSession = s.activeSSHSession()
+	s.DefaultUserName = s.getDefaultUser()
+	segmentTemplate := s.props.getString(SegmentTemplate, "")
+	if segmentTemplate != "" {
+		s.Root = s.env.isRunningAsRoot()
+		template := &textTemplate{
+			Template: segmentTemplate,
+			Context:  s,
+		}
+		s.templateText = template.render()
+		return len(s.templateText) > 0
+	}
 	showDefaultUser := s.props.getBool(DisplayDefault, true)
-	defaultUser := s.props.getString(DefaultUserName, "")
 	if !showDefaultUser && defaultUser == s.UserName {
 		return false
 	}
@@ -51,16 +67,8 @@ func (s *session) init(props *properties, env environmentInfo) {
 }
 
 func (s *session) getFormattedText() string {
-	s.ComputerName = s.getComputerName()
-	s.SSHSession = s.activeSSHSession()
-	segmentTemplate := s.props.getString(SegmentTemplate, "")
-	if segmentTemplate != "" {
-		s.Root = s.env.isRunningAsRoot()
-		template := &textTemplate{
-			Template: segmentTemplate,
-			Context:  s,
-		}
-		return template.render()
+	if len(s.templateText) > 0 {
+		return s.templateText
 	}
 	separator := ""
 	if s.props.getBool(DisplayHost, true) && s.props.getBool(DisplayUser, true) {
@@ -98,6 +106,14 @@ func (s *session) getUserName() string {
 	return username
 }
 
+func (s *session) getDefaultUser() string {
+	user := s.env.getenv(defaultUserEnvVar)
+	if len(user) == 0 {
+		user = s.props.getString(DefaultUserName, "")
+	}
+	return user
+}
+
 func (s *session) activeSSHSession() bool {
 	keys := []string{
 		"SSH_CONNECTION",

diff --git a/src/segment_session_test.go b/src/segment_session_test.go
@@ -14,6 +14,7 @@ type sessionArgs struct {
 	connection        string
 	client            string
 	sshIcon           string
+	defaultUserName   string
 }
 
 func setupSession(args *sessionArgs) session {
@@ -23,6 +24,7 @@ func setupSession(args *sessionArgs) session {
 	env.On("getRuntimeGOOS", nil).Return(args.goos)
 	env.On("getenv", "SSH_CONNECTION").Return(args.connection)
 	env.On("getenv", "SSH_CLIENT").Return(args.client)
+	env.On("getenv", defaultUserEnvVar).Return(args.defaultUserName)
 	props := &properties{
 		values: map[Property]interface{}{
 			UserInfoSeparator: args.userInfoSeparator,
@@ -154,59 +156,89 @@ func TestActiveSSHSessionActiveBoth(t *testing.T) {
 
 func TestSessionSegmentTemplate(t *testing.T) {
 	cases := []struct {
-		Case           string
-		ExpectedString string
-		UserName       string
-		ComputerName   string
-		SSHSession     bool
-		Root           bool
-		Template       string
+		Case            string
+		ExpectedEnabled bool
+		ExpectedString  string
+		UserName        string
+		DefaultUserName string
+		ComputerName    string
+		SSHSession      bool
+		Root            bool
+		Template        string
 	}{
 		{
-			Case:           "user and computer",
-			ExpectedString: "john@company-laptop",
-			ComputerName:   "company-laptop",
-			UserName:       "john",
-			Template:       "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+			Case:            "user and computer",
+			ExpectedString:  "john@company-laptop",
+			ComputerName:    "company-laptop",
+			UserName:        "john",
+			Template:        "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+			ExpectedEnabled: true,
 		},
 		{
-			Case:           "user only",
-			ExpectedString: "john",
-			UserName:       "john",
-			Template:       "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+			Case:            "user only",
+			ExpectedString:  "john",
+			UserName:        "john",
+			Template:        "{{.UserName}}{{if .ComputerName}}@{{.ComputerName}}{{end}}",
+			ExpectedEnabled: true,
 		},
 		{
-			Case:           "user with ssh",
-			ExpectedString: "john on remote",
-			UserName:       "john",
-			SSHSession:     true,
-			ComputerName:   "remote",
-			Template:       "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+			Case:            "user with ssh",
+			ExpectedString:  "john on remote",
+			UserName:        "john",
+			SSHSession:      true,
+			ComputerName:    "remote",
+			Template:        "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+			ExpectedEnabled: true,
 		},
 		{
-			Case:           "user without ssh",
-			ExpectedString: "john",
-			UserName:       "john",
-			SSHSession:     false,
-			ComputerName:   "remote",
-			Template:       "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+			Case:            "user without ssh",
+			ExpectedString:  "john",
+			UserName:        "john",
+			SSHSession:      false,
+			ComputerName:    "remote",
+			Template:        "{{.UserName}}{{if .SSHSession}} on {{.ComputerName}}{{end}}",
+			ExpectedEnabled: true,
 		},
 		{
-			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}}",
+			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,
 		},
 		{
-			Case:           "no template",
-			ExpectedString: "\uf817 <>john</>@<>remote</>",
-			UserName:       "john",
-			SSHSession:     true,
-			ComputerName:   "remote",
-			Root:           true,
+			Case:            "no template",
+			ExpectedString:  "\uf817 <>john</>@<>remote</>",
+			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,
 		},
 	}
 
@@ -222,6 +254,7 @@ func TestSessionSegmentTemplate(t *testing.T) {
 		env.On("getenv", "SSH_CONNECTION").Return(SSHSession)
 		env.On("getenv", "SSH_CLIENT").Return(SSHSession)
 		env.On("isRunningAsRoot", nil).Return(tc.Root)
+		env.On("getenv", defaultUserEnvVar).Return(tc.DefaultUserName)
 		props := &properties{
 			values: map[Property]interface{}{
 				SegmentTemplate: tc.Template,
@@ -231,7 +264,9 @@ func TestSessionSegmentTemplate(t *testing.T) {
 			env:   env,
 			props: props,
 		}
-		_ = session.enabled()
-		assert.Equal(t, tc.ExpectedString, session.string(), tc.Case)
+		assert.Equal(t, tc.ExpectedEnabled, session.enabled(), tc.Case)
+		if tc.ExpectedEnabled {
+			assert.Equal(t, tc.ExpectedString, session.string(), tc.Case)
+		}
 	}
 }