From fb57ef18b6837d2eb553e2e2979249f577f9c388 Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Thu, 22 Apr 2021 20:18:58 +0200
Subject: [PATCH] feat: display message on az update needed

resolves #669
---
 src/segment_az.go      | 11 +++++++
 src/segment_az_test.go | 70 +++++++++++++++++++++++++-----------------
 2 files changed, 52 insertions(+), 29 deletions(-)

diff --git a/src/segment_az.go b/src/segment_az.go
index a51f3846..4b2afb5e 100644
--- a/src/segment_az.go
+++ b/src/segment_az.go
@@ -19,6 +19,11 @@ const (
 	DisplaySubscriptionID Property = "display_id"
 	// DisplaySubscriptionName hides or shows the subscription display name
 	DisplaySubscriptionName Property = "display_name"
+
+	updateConsentNeeded = "Do you want to continue?"
+	updateMessage       = "AZ CLI: Update needed!"
+	updateForeground    = "#ffffff"
+	updateBackground    = "#ff5349"
 )
 
 func (a *az) string() string {
@@ -68,6 +73,12 @@ func (a *az) getFromAzCli() (string, string, bool) {
 		return "", "", false
 	}
 
+	if strings.Contains(output, updateConsentNeeded) {
+		a.props.foreground = updateForeground
+		a.props.background = updateBackground
+		return updateMessage, "", true
+	}
+
 	splittedOutput := strings.Split(output, "\n")
 	if len(splittedOutput) < 2 {
 		return "", "", false
diff --git a/src/segment_az_test.go b/src/segment_az_test.go
index 0df232aa..c9656744 100644
--- a/src/segment_az_test.go
+++ b/src/segment_az_test.go
@@ -21,7 +21,8 @@ func TestAzSegment(t *testing.T) {
 		DisplayID       bool
 		DisplayName     bool
 	}{
-		{Case: "envvars present",
+		{
+			Case:            "envvars present",
 			ExpectedEnabled: true,
 			ExpectedString:  "foo$bar",
 			EnvSubName:      "foo",
@@ -29,78 +30,89 @@ func TestAzSegment(t *testing.T) {
 			CliExists:       false,
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     true},
-		{Case: "envvar name present",
+			DisplayName:     true,
+		},
+		{
+			Case:            "envvar name present",
 			ExpectedEnabled: true,
 			ExpectedString:  "foo$",
 			EnvSubName:      "foo",
-			EnvSubID:        "",
 			CliExists:       false,
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     true},
-		{Case: "envvar id present",
+			DisplayName:     true,
+		},
+		{
+			Case:            "envvar id present",
 			ExpectedEnabled: true,
 			ExpectedString:  "$bar",
-			EnvSubName:      "",
 			EnvSubID:        "bar",
 			CliExists:       false,
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     true},
-		{Case: "cli not found",
+			DisplayName:     true,
+		},
+		{
+			Case:            "cli not found",
 			ExpectedEnabled: false,
 			ExpectedString:  "$",
-			EnvSubName:      "",
-			EnvSubID:        "",
 			CliExists:       false,
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     true},
-		{Case: "cli contains data",
+			DisplayName:     true,
+		},
+		{
+			Case:            "cli contains data",
 			ExpectedEnabled: true,
 			ExpectedString:  "foo$bar",
-			EnvSubName:      "",
-			EnvSubID:        "",
 			CliExists:       true,
 			CliSubName:      "foo",
 			CliSubID:        "bar",
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     true},
-		{Case: "print only name",
+			DisplayName:     true,
+		},
+		{
+			Case:            "print only name",
 			ExpectedEnabled: true,
 			ExpectedString:  "foo",
-			EnvSubName:      "",
-			EnvSubID:        "",
 			CliExists:       true,
 			CliSubName:      "foo",
 			CliSubID:        "bar",
 			InfoSeparator:   "$",
 			DisplayID:       false,
-			DisplayName:     true},
-		{Case: "print only id",
+			DisplayName:     true,
+		},
+		{
+			Case:            "print only id",
 			ExpectedEnabled: true,
 			ExpectedString:  "bar",
-			EnvSubName:      "",
-			EnvSubID:        "",
 			CliExists:       true,
 			CliSubName:      "foo",
 			CliSubID:        "bar",
 			InfoSeparator:   "$",
 			DisplayID:       true,
-			DisplayName:     false},
-		{Case: "print none",
+			DisplayName:     false,
+		},
+		{
+			Case:            "print none",
 			ExpectedEnabled: false,
-			ExpectedString:  "",
-			EnvSubName:      "",
-			EnvSubID:        "",
 			CliExists:       true,
 			CliSubName:      "foo",
 			CliSubID:        "bar",
 			InfoSeparator:   "$",
 			DisplayID:       false,
-			DisplayName:     false},
+			DisplayName:     false,
+		},
+		{
+			Case:            "update needed",
+			ExpectedEnabled: true,
+			ExpectedString:  updateMessage,
+			CliExists:       true,
+			CliSubName:      "Do you want to continue? (Y/n): Visual Studio Enterprise",
+			DisplayID:       false,
+			DisplayName:     true,
+		},
 	}
 
 	for _, tc := range cases {