From 8f991cb9f8429820218f0543642d9d33f3d9f5c0 Mon Sep 17 00:00:00 2001
From: Jack Francis <jack.francis@microsoft.com>
Date: Thu, 11 Apr 2019 13:42:00 -0700
Subject: [PATCH] chore: more tests, get rid of dead code, rationalize
 GetPrimaryAvailabilitySetName

---
 pkg/api/types.go                 |  4 ++-
 pkg/api/types_test.go            | 44 +++++++++++++++++++++++++++++++-
 pkg/engine/armvariables.go       |  3 +--
 pkg/engine/template_generator.go |  6 -----
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/pkg/api/types.go b/pkg/api/types.go
index 0ecf5c794f1..39a64c7ff42 100644
--- a/pkg/api/types.go
+++ b/pkg/api/types.go
@@ -889,7 +889,9 @@ func (p *Properties) GetNSGName() string {
 // GetPrimaryAvailabilitySetName returns the name of the primary availability set of the cluster
 func (p *Properties) GetPrimaryAvailabilitySetName() string {
 	if p.AgentPoolProfiles != nil {
-		return p.AgentPoolProfiles[0].Name + "-availabilitySet-" + p.GetClusterID()
+		if p.AgentPoolProfiles[0].AvailabilityProfile == AvailabilitySet {
+			return p.AgentPoolProfiles[0].Name + "-availabilitySet-" + p.GetClusterID()
+		}
 	}
 	return ""
 }
diff --git a/pkg/api/types_test.go b/pkg/api/types_test.go
index 606b5013450..c14936ae668 100644
--- a/pkg/api/types_test.go
+++ b/pkg/api/types_test.go
@@ -2207,6 +2207,27 @@ func TestGetPrimaryAvailabilitySetName(t *testing.T) {
 	if got != expected {
 		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
 	}
+
+	p.AgentPoolProfiles = []*AgentPoolProfile{
+		{
+			Name:                "agentpool",
+			VMSize:              "Standard_D2_v2",
+			Count:               1,
+			AvailabilityProfile: VirtualMachineScaleSets,
+		},
+	}
+	expected = ""
+	got = p.GetPrimaryAvailabilitySetName()
+	if got != expected {
+		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
+	}
+
+	p.AgentPoolProfiles = nil
+	expected = ""
+	got = p.GetPrimaryAvailabilitySetName()
+	if got != expected {
+		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
+	}
 }
 
 func TestGetPrimaryScaleSetName(t *testing.T) {
@@ -2232,7 +2253,7 @@ func TestGetPrimaryScaleSetName(t *testing.T) {
 	expected := "k8s-agentpool-28513887-vmss"
 	got := p.GetPrimaryScaleSetName()
 	if got != expected {
-		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
+		t.Errorf("expected primary scaleset name %s, but got %s", expected, got)
 	}
 
 	// Test with Windows agentpool
@@ -2240,6 +2261,27 @@ func TestGetPrimaryScaleSetName(t *testing.T) {
 
 	expected = "2851k8s00"
 	got = p.GetPrimaryScaleSetName()
+	if got != expected {
+		t.Errorf("expected primary scaleset name %s, but got %s", expected, got)
+	}
+
+	p.AgentPoolProfiles = []*AgentPoolProfile{
+		{
+			Name:                "agentpool",
+			VMSize:              "Standard_D2_v2",
+			Count:               1,
+			AvailabilityProfile: AvailabilitySet,
+		},
+	}
+	expected = ""
+	got = p.GetPrimaryScaleSetName()
+	if got != expected {
+		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
+	}
+
+	p.AgentPoolProfiles = nil
+	expected = ""
+	got = p.GetPrimaryScaleSetName()
 	if got != expected {
 		t.Errorf("expected primary availability set name %s, but got %s", expected, got)
 	}
diff --git a/pkg/engine/armvariables.go b/pkg/engine/armvariables.go
index cf3e84c8ad2..146fd24acc9 100644
--- a/pkg/engine/armvariables.go
+++ b/pkg/engine/armvariables.go
@@ -326,7 +326,6 @@ func getK8sMasterVars(cs *api.ContainerService) (map[string]interface{}, error)
 	}
 
 	if cs.Properties.AnyAgentUsesVirtualMachineScaleSets() {
-		masterVars["primaryScaleSetName"] = cs.Properties.GetPrimaryScaleSetName()
 		masterVars["primaryAvailabilitySetName"] = ""
 		masterVars["vmType"] = "vmss"
 	} else {
@@ -335,9 +334,9 @@ func getK8sMasterVars(cs *api.ContainerService) (map[string]interface{}, error)
 		} else {
 			masterVars["primaryAvailabilitySetName"] = ""
 		}
-		masterVars["primaryScaleSetName"] = ""
 		masterVars["vmType"] = "standard"
 	}
+	masterVars["primaryScaleSetName"] = cs.Properties.GetPrimaryScaleSetName()
 
 	if isHostedMaster {
 		masterVars["kubernetesAPIServerIP"] = "[parameters('kubernetesEndpoint')]"
diff --git a/pkg/engine/template_generator.go b/pkg/engine/template_generator.go
index 18e7b720513..27719088233 100644
--- a/pkg/engine/template_generator.go
+++ b/pkg/engine/template_generator.go
@@ -968,12 +968,6 @@ func (t *TemplateGenerator) getTemplateFuncMap(cs *api.ContainerService) templat
 		"GetMasterEtcdClientPort": func() int {
 			return DefaultMasterEtcdClientPort
 		},
-		"GetPrimaryAvailabilitySetName": func() string {
-			return cs.Properties.GetPrimaryAvailabilitySetName()
-		},
-		"GetPrimaryScaleSetName": func() string {
-			return cs.Properties.GetPrimaryScaleSetName()
-		},
 		"UseCloudControllerManager": func() bool {
 			return cs.Properties.OrchestratorProfile.KubernetesConfig.UseCloudControllerManager != nil && *cs.Properties.OrchestratorProfile.KubernetesConfig.UseCloudControllerManager
 		},