From fc2e1c6b08b46ab4d5360cd9430a95daf6d92d70 Mon Sep 17 00:00:00 2001
From: Ahmet Alp Balkan <ahmetb@google.com>
Date: Sat, 18 Apr 2020 16:17:47 -0700
Subject: [PATCH] Fix bug about where cur ns was stored in yaml

Signed-off-by: Ahmet Alp Balkan <ahmetb@google.com>
---
 internal/kubeconfig/namespace.go       | 29 ++++++++++++++++++++++----
 internal/testutil/kubeconfigbuilder.go |  8 ++++---
 internal/testutil/kubeconfigloader.go  |  1 -
 test/kubens.bats                       |  1 -
 4 files changed, 30 insertions(+), 9 deletions(-)
 delete mode 100644 internal/testutil/kubeconfigloader.go

diff --git a/internal/kubeconfig/namespace.go b/internal/kubeconfig/namespace.go
index fa9921e..f8e0086 100644
--- a/internal/kubeconfig/namespace.go
+++ b/internal/kubeconfig/namespace.go
@@ -11,7 +11,11 @@ func (k *Kubeconfig) NamespaceOfContext(contextName string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	ns := valueOf(ctx, "namespace")
+	ctxBody := valueOf(ctx, "context")
+	if ctxBody == nil {
+		return defaultNamespace, nil
+	}
+	ns := valueOf(ctxBody, "namespace")
 	if ns == nil || ns.Value == "" {
 		return defaultNamespace, nil
 	}
@@ -19,11 +23,21 @@ func (k *Kubeconfig) NamespaceOfContext(contextName string) (string, error) {
 }
 
 func (k *Kubeconfig) SetNamespace(ctxName string, ns string) error {
-	ctx, err := k.contextNode(ctxName)
+	ctxNode, err := k.contextNode(ctxName)
 	if err != nil {
 		return err
 	}
-	nsNode := valueOf(ctx, "namespace")
+
+	var ctxBodyNodeWasEmpty bool // actual namespace value is in contexts[index].context.namespace, but .context might not exist
+	ctxBodyNode := valueOf(ctxNode, "context")
+	if ctxBodyNode == nil {
+		ctxBodyNodeWasEmpty = true
+		ctxBodyNode = &yaml.Node{
+			Kind: yaml.MappingNode,
+		}
+	}
+
+	nsNode := valueOf(ctxBodyNode, "namespace")
 	if nsNode != nil {
 		nsNode.Value = ns
 		return nil
@@ -37,6 +51,13 @@ func (k *Kubeconfig) SetNamespace(ctxName string, ns string) error {
 		Kind:  yaml.ScalarNode,
 		Value: ns,
 		Tag:   "!!str"}
-	ctx.Content = append(ctx.Content, keyNode, valueNode)
+	ctxBodyNode.Content = append(ctxBodyNode.Content, keyNode, valueNode)
+	if ctxBodyNodeWasEmpty {
+		ctxNode.Content = append(ctxNode.Content, &yaml.Node{
+			Kind:  yaml.ScalarNode,
+			Value: "context",
+			Tag:   "!!str",
+		}, ctxBodyNode)
+	}
 	return nil
 }
diff --git a/internal/testutil/kubeconfigbuilder.go b/internal/testutil/kubeconfigbuilder.go
index 8fdb0d9..b046ff3 100644
--- a/internal/testutil/kubeconfigbuilder.go
+++ b/internal/testutil/kubeconfigbuilder.go
@@ -8,12 +8,14 @@ import (
 )
 
 type Context struct {
-	Name      string `yaml:"name,omitempty"`
-	Namespace string `yaml:"namespace,omitempty"`
+	Name    string `yaml:"name,omitempty"`
+	Context struct {
+		Namespace string `yaml:"namespace,omitempty"`
+	} `yaml:"context,omitempty"`
 }
 
 func Ctx(name string) *Context           { return &Context{Name: name} }
-func (c *Context) Ns(ns string) *Context { c.Namespace = ns; return c }
+func (c *Context) Ns(ns string) *Context { c.Context.Namespace = ns; return c }
 
 type Kubeconfig map[string]interface{}
 
diff --git a/internal/testutil/kubeconfigloader.go b/internal/testutil/kubeconfigloader.go
deleted file mode 100644
index 110b2e6..0000000
--- a/internal/testutil/kubeconfigloader.go
+++ /dev/null
@@ -1 +0,0 @@
-package testutil
diff --git a/test/kubens.bats b/test/kubens.bats
index 6bbcb09..847f4c2 100644
--- a/test/kubens.bats
+++ b/test/kubens.bats
@@ -21,7 +21,6 @@ load common
   run ${COMMAND}
   echo "$output"
   [[ "$status" -eq 1 ]]
-  [[ "$output" = *"current-context is not set"* ]]
 }
 
 @test "list namespaces" {