From a1d630015e7ce0b84aa056acbd593e7dedf8bcbe Mon Sep 17 00:00:00 2001
From: Fred Lotter <fred.lotter@canonical.com>
Date: Wed, 28 Aug 2024 21:21:31 +0200
Subject: [PATCH] simplify extension registration in tests

---
 internals/overlord/planstate/manager_test.go | 12 +++++++++---
 internals/overlord/planstate/package_test.go |  8 +++-----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/internals/overlord/planstate/manager_test.go b/internals/overlord/planstate/manager_test.go
index 567d4831..ef28f43d 100644
--- a/internals/overlord/planstate/manager_test.go
+++ b/internals/overlord/planstate/manager_test.go
@@ -110,10 +110,12 @@ var extLoadLayers = []string{`
 `}
 
 func (ps *planSuite) TestExtLoadLayers(c *C) {
+	plan.RegisterExtension(testField, testExtension{})
+	defer plan.UnregisterExtension(testField)
+
 	var err error
 	ps.planMgr, err = planstate.NewManager(ps.layersDir)
 	c.Assert(err, IsNil)
-	plan.RegisterExtension(testField, testExtension{})
 	// Write layers
 	for _, l := range extLoadLayers {
 		ps.writeLayer(c, string(reindent(l)))
@@ -224,10 +226,12 @@ services:
 }
 
 func (ps *planSuite) TestExtAppendLayers(c *C) {
+	plan.RegisterExtension(testField, testExtension{})
+	defer plan.UnregisterExtension(testField)
+
 	var err error
 	ps.planMgr, err = planstate.NewManager(ps.layersDir)
 	c.Assert(err, IsNil)
-	plan.RegisterExtension(testField, testExtension{})
 
 	// Append a layer when there are no layers.
 	layer := ps.parseLayer(c, 0, "label1", `
@@ -462,10 +466,12 @@ checks:
 }
 
 func (ps *planSuite) TestExtCombineLayers(c *C) {
+	plan.RegisterExtension(testField, testExtension{})
+	defer plan.UnregisterExtension(testField)
+
 	var err error
 	ps.planMgr, err = planstate.NewManager(ps.layersDir)
 	c.Assert(err, IsNil)
-	plan.RegisterExtension(testField, testExtension{})
 
 	// "Combine" layer with no layers should just append.
 	layer := ps.parseLayer(c, 0, "label1", `
diff --git a/internals/overlord/planstate/package_test.go b/internals/overlord/planstate/package_test.go
index bfb0d765..4e79d7e6 100644
--- a/internals/overlord/planstate/package_test.go
+++ b/internals/overlord/planstate/package_test.go
@@ -42,8 +42,6 @@ type planSuite struct {
 var _ = Suite(&planSuite{})
 
 func (ps *planSuite) SetUpTest(c *C) {
-	plan.UnregisterExtension(testField)
-
 	ps.layersDir = filepath.Join(c.MkDir(), "layers")
 	err := os.Mkdir(ps.layersDir, 0755)
 	c.Assert(err, IsNil)
@@ -110,7 +108,7 @@ const testField string = "test-field"
 // testExtension implements the LayerSectionExtension interface.
 type testExtension struct{}
 
-func (te testExtension) ParseSection(data yaml.Node) (plan.LayerSection, error) {
+func (te testExtension) ParseSection(data yaml.Node) (plan.Section, error) {
 	ts := &testSection{}
 	err := data.Decode(ts)
 	if err != nil {
@@ -125,7 +123,7 @@ func (te testExtension) ParseSection(data yaml.Node) (plan.LayerSection, error)
 	return ts, nil
 }
 
-func (te testExtension) CombineSections(sections ...plan.LayerSection) (plan.LayerSection, error) {
+func (te testExtension) CombineSections(sections ...plan.Section) (plan.Section, error) {
 	ts := &testSection{}
 	for _, section := range sections {
 		err := ts.Combine(section)
@@ -161,7 +159,7 @@ func (ts *testSection) IsZero() bool {
 	return ts.Entries == nil
 }
 
-func (ts *testSection) Combine(other plan.LayerSection) error {
+func (ts *testSection) Combine(other plan.Section) error {
 	otherTSection, ok := other.(*testSection)
 	if !ok {
 		return fmt.Errorf("invalid section type")