From dabc9d4f209285050140b093bb5374e10e8a514c Mon Sep 17 00:00:00 2001
From: nwykes <nwykes@gmail.com>
Date: Wed, 7 Oct 2020 09:14:26 -0600
Subject: [PATCH] feat: session segment updates

only display UserInfoSeparator if both username and hostname are dislayed
cleanHostName only if it ends with garbage instead of contains.
  fixes .localdomain having just .local stripped
add .localdomain to garbage
---
 environment.go          |  5 ++++-
 segment_session.go      |  6 +++++-
 segment_session_test.go | 25 ++++++++++++++++++++++++-
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/environment.go b/environment.go
index f7460acf..a8d96d31 100755
--- a/environment.go
+++ b/environment.go
@@ -138,9 +138,12 @@ func cleanHostName(hostName string) string {
 	garbage := []string{
 		".lan",
 		".local",
+		".localdomain",
 	}
 	for _, g := range garbage {
-		hostName = strings.Replace(hostName, g, "", 1)
+		if strings.HasSuffix(hostName, g) {
+			hostName = strings.Replace(hostName, g, "", 1)
+		}
 	}
 	return hostName
 }
diff --git a/segment_session.go b/segment_session.go
index f7cc4f89..f3fc3870 100755
--- a/segment_session.go
+++ b/segment_session.go
@@ -39,7 +39,11 @@ func (s *session) init(props *properties, env environmentInfo) {
 func (s *session) getFormattedText() string {
 	username := s.getUserName()
 	computername := s.getComputerName()
-	return fmt.Sprintf("<%s>%s</>%s<%s>%s</>", s.props.getColor(UserColor, s.props.foreground), username, s.props.getString(UserInfoSeparator, "@"), s.props.getColor(HostColor, s.props.foreground), computername)
+	separator := ""
+	if s.props.getBool(DisplayHost, true) && s.props.getBool(DisplayUser, true) {
+		separator = s.props.getString(UserInfoSeparator, "@")
+	}
+	return fmt.Sprintf("<%s>%s</>%s<%s>%s</>", s.props.getColor(UserColor, s.props.foreground), username, separator, s.props.getColor(HostColor, s.props.foreground), computername)
 }
 
 func (s *session) getComputerName() string {
diff --git a/segment_session_test.go b/segment_session_test.go
index 9c05ac41..bf6423f9 100755
--- a/segment_session_test.go
+++ b/segment_session_test.go
@@ -7,7 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func testUserInfoWriter(userInfoSeparator string, username string, hostname string, goos string) string {
+func setupSession(userInfoSeparator string, username string, hostname string, goos string) session {
 	env := new(MockedEnvironment)
 	user := user.User{
 		Username: username,
@@ -24,6 +24,11 @@ func testUserInfoWriter(userInfoSeparator string, username string, hostname stri
 		env:   env,
 		props: props,
 	}
+	return s
+}
+
+func testUserInfoWriter(userInfoSeparator string, username string, hostname string, goos string) string {
+	s := setupSession(userInfoSeparator, username, hostname, goos)
 	return s.getFormattedText()
 }
 
@@ -38,3 +43,21 @@ func TestWriteUserInfoWindowsIncludingHostname(t *testing.T) {
 	got := testUserInfoWriter("@", "surface\\bill", "surface", "windows")
 	assert.EqualValues(t, want, got)
 }
+
+func TestWriteOnlyUsername(t *testing.T) {
+	s := setupSession("@", "surface\\bill", "surface", "windows")
+	s.props.values[DisplayHost] = false
+
+	want := "<#fff>bill</><#fff></>"
+	got := s.getFormattedText()
+	assert.EqualValues(t, want, got)
+}
+
+func TestWriteOnlyHostname(t *testing.T) {
+	s := setupSession("@", "surface\\bill", "surface", "windows")
+	s.props.values[DisplayUser] = false
+
+	want := "<#fff></><#fff>surface</>"
+	got := s.getFormattedText()
+	assert.EqualValues(t, want, got)
+}