From 744a8820fb6b8e499845b9cb60c2569560c1e11d Mon Sep 17 00:00:00 2001
From: Tom Clegg <tom@tomclegg.ca>
Date: Wed, 1 Apr 2020 16:44:02 -0400
Subject: [PATCH] 16270: Fill in missing scratch fields on InstanceType
 entries.

Previously they were being filled in correctly when written as an
array in the config file, but not when written as a map.

Arvados-DCO-1.1-Signed-off-by: Tom Clegg <tom@tomclegg.ca>
---
 sdk/go/arvados/config.go      | 39 +++++++++++++++++++++--------------
 sdk/go/arvados/config_test.go | 26 +++++++++++++++++++++++
 2 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/sdk/go/arvados/config.go b/sdk/go/arvados/config.go
index a70980cbde..79e47ba5d1 100644
--- a/sdk/go/arvados/config.go
+++ b/sdk/go/arvados/config.go
@@ -421,6 +421,24 @@ var errDuplicateInstanceTypeName = errors.New("duplicate instance type name")
 // UnmarshalJSON handles old config files that provide an array of
 // instance types instead of a hash.
 func (it *InstanceTypeMap) UnmarshalJSON(data []byte) error {
+	fixup := func(t InstanceType) (InstanceType, error) {
+		if t.ProviderType == "" {
+			t.ProviderType = t.Name
+		}
+		if t.Scratch == 0 {
+			t.Scratch = t.IncludedScratch + t.AddedScratch
+		} else if t.AddedScratch == 0 {
+			t.AddedScratch = t.Scratch - t.IncludedScratch
+		} else if t.IncludedScratch == 0 {
+			t.IncludedScratch = t.Scratch - t.AddedScratch
+		}
+
+		if t.Scratch != (t.IncludedScratch + t.AddedScratch) {
+			return t, fmt.Errorf("InstanceType %q: Scratch != (IncludedScratch + AddedScratch)", t.Name)
+		}
+		return t, nil
+	}
+
 	if len(data) > 0 && data[0] == '[' {
 		var arr []InstanceType
 		err := json.Unmarshal(data, &arr)
@@ -436,19 +454,9 @@ func (it *InstanceTypeMap) UnmarshalJSON(data []byte) error {
 			if _, ok := (*it)[t.Name]; ok {
 				return errDuplicateInstanceTypeName
 			}
-			if t.ProviderType == "" {
-				t.ProviderType = t.Name
-			}
-			if t.Scratch == 0 {
-				t.Scratch = t.IncludedScratch + t.AddedScratch
-			} else if t.AddedScratch == 0 {
-				t.AddedScratch = t.Scratch - t.IncludedScratch
-			} else if t.IncludedScratch == 0 {
-				t.IncludedScratch = t.Scratch - t.AddedScratch
-			}
-
-			if t.Scratch != (t.IncludedScratch + t.AddedScratch) {
-				return fmt.Errorf("%v: Scratch != (IncludedScratch + AddedScratch)", t.Name)
+			t, err := fixup(t)
+			if err != nil {
+				return err
 			}
 			(*it)[t.Name] = t
 		}
@@ -464,8 +472,9 @@ func (it *InstanceTypeMap) UnmarshalJSON(data []byte) error {
 	*it = InstanceTypeMap(hash)
 	for name, t := range *it {
 		t.Name = name
-		if t.ProviderType == "" {
-			t.ProviderType = name
+		t, err := fixup(t)
+		if err != nil {
+			return err
 		}
 		(*it)[name] = t
 	}
diff --git a/sdk/go/arvados/config_test.go b/sdk/go/arvados/config_test.go
index b984cb5669..e4d26e03fd 100644
--- a/sdk/go/arvados/config_test.go
+++ b/sdk/go/arvados/config_test.go
@@ -45,3 +45,29 @@ func (s *ConfigSuite) TestInstanceTypeSize(c *check.C) {
 	c.Check(int64(it.Scratch), check.Equals, int64(4000000000))
 	c.Check(int64(it.RAM), check.Equals, int64(4294967296))
 }
+
+func (s *ConfigSuite) TestInstanceTypeFixup(c *check.C) {
+	for _, confdata := range []string{
+		// Current format: map of entries
+		`{foo4: {IncludedScratch: 4GB}, foo8: {ProviderType: foo_8, Scratch: 8GB}}`,
+		// Legacy format: array of entries with key in "Name" field
+		`[{Name: foo4, IncludedScratch: 4GB}, {Name: foo8, ProviderType: foo_8, Scratch: 8GB}]`,
+	} {
+		c.Log(confdata)
+		var itm InstanceTypeMap
+		err := yaml.Unmarshal([]byte(confdata), &itm)
+		c.Check(err, check.IsNil)
+
+		c.Check(itm["foo4"].Name, check.Equals, "foo4")
+		c.Check(itm["foo4"].ProviderType, check.Equals, "foo4")
+		c.Check(itm["foo4"].Scratch, check.Equals, ByteSize(4000000000))
+		c.Check(itm["foo4"].AddedScratch, check.Equals, ByteSize(0))
+		c.Check(itm["foo4"].IncludedScratch, check.Equals, ByteSize(4000000000))
+
+		c.Check(itm["foo8"].Name, check.Equals, "foo8")
+		c.Check(itm["foo8"].ProviderType, check.Equals, "foo_8")
+		c.Check(itm["foo8"].Scratch, check.Equals, ByteSize(8000000000))
+		c.Check(itm["foo8"].AddedScratch, check.Equals, ByteSize(8000000000))
+		c.Check(itm["foo8"].IncludedScratch, check.Equals, ByteSize(0))
+	}
+}
-- 
2.30.2