From 157c4eb86b1a4fbbf798fa36355ff802cd39b4a3 Mon Sep 17 00:00:00 2001
From: Bevan Arps <bevan.arps@microsoft.com>
Date: Thu, 1 Jul 2021 15:50:22 +1200
Subject: [PATCH] Integrate into pipeline

---
 hack/generator/pkg/codegen/code_generator.go                 | 1 +
 hack/generator/pkg/codegen/golden_files_test.go              | 1 +
 hack/generator/pkg/codegen/pipeline/mark_storage_version.go  | 5 ++++-
 .../pkg/codegen/testdata/ARMCodeGeneratorPipeline.golden     | 1 +
 4 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/hack/generator/pkg/codegen/code_generator.go b/hack/generator/pkg/codegen/code_generator.go
index 558626125aa..086b56b720d 100644
--- a/hack/generator/pkg/codegen/code_generator.go
+++ b/hack/generator/pkg/codegen/code_generator.go
@@ -168,6 +168,7 @@ func createAllPipelineStages(idFactory astmodel.IdentifierFactory, configuration
 		pipeline.InjectJsonSerializationTests(idFactory).UsedFor(pipeline.ARMTarget),
 
 		pipeline.MarkStorageVersion(),
+		pipeline.InjectHubFunction(idFactory).UsedFor(pipeline.ARMTarget),
 
 		// Safety checks at the end:
 		pipeline.EnsureDefinitionsDoNotUseAnyTypes(),
diff --git a/hack/generator/pkg/codegen/golden_files_test.go b/hack/generator/pkg/codegen/golden_files_test.go
index b39d18d63d0..17c8cc60291 100644
--- a/hack/generator/pkg/codegen/golden_files_test.go
+++ b/hack/generator/pkg/codegen/golden_files_test.go
@@ -153,6 +153,7 @@ func NewTestCodeGenerator(testName string, path string, t *testing.T, testConfig
 			"injectOriginalVersionFunction",
 			"injectOriginalVersionProperty",
 			"injectPropertyAssignmentFunctions",
+			pipeline.InjectHubFunctionStageId,
 			"reportTypesAndVersions")
 		if !testConfig.HasARMResources {
 			codegen.RemoveStages("createArmTypes", "applyArmConversionInterface")
diff --git a/hack/generator/pkg/codegen/pipeline/mark_storage_version.go b/hack/generator/pkg/codegen/pipeline/mark_storage_version.go
index e24b94954d7..45c92fd3d8a 100644
--- a/hack/generator/pkg/codegen/pipeline/mark_storage_version.go
+++ b/hack/generator/pkg/codegen/pipeline/mark_storage_version.go
@@ -14,10 +14,13 @@ import (
 	"github.com/Azure/azure-service-operator/hack/generator/pkg/astmodel"
 )
 
+// MarkStorageVersionStageId is the unique identifier for this pipeline stage
+const MarkStorageVersionStageId = "markStorageVersion"
+
 // MarkStorageVersion creates a Stage to mark a particular version as a storage version
 func MarkStorageVersion() Stage {
 	return MakeStage(
-		"markStorageVersion",
+		MarkStorageVersionStageId,
 		"Mark the latest version of each resource as the storage version",
 		func(ctx context.Context, types astmodel.Types) (astmodel.Types, error) {
 			updatedDefs, err := MarkLatestResourceVersionsForStorage(types)
diff --git a/hack/generator/pkg/codegen/testdata/ARMCodeGeneratorPipeline.golden b/hack/generator/pkg/codegen/testdata/ARMCodeGeneratorPipeline.golden
index 8c67162bb07..000275ea56e 100644
--- a/hack/generator/pkg/codegen/testdata/ARMCodeGeneratorPipeline.golden
+++ b/hack/generator/pkg/codegen/testdata/ARMCodeGeneratorPipeline.golden
@@ -37,6 +37,7 @@ injectOriginalGVKFunction           azure      Inject the function OriginalGVK()
 simplifyDefinitions                            Flatten definitions by removing wrapper types
 jsonTestCases                       azure      Add test cases to verify JSON serialization
 markStorageVersion                             Mark the latest version of each resource as the storage version
+injectHubFunction                   azure      Inject the function Hub() into each hub resource
 rogueCheck                                     Check for rogue definitions using AnyTypes
 ensureArmTypeExistsForEveryType     azure      Check that an ARM type exists for both Spec and Status of each resource
 deleteGenerated                                Delete generated code from .