From 7abbcd21bd2a9b40de49a8d0a9003853f0502216 Mon Sep 17 00:00:00 2001
From: Sebastian Czech <sczech@paloaltonetworks.com>
Date: Wed, 3 Apr 2024 15:46:46 +0200
Subject: [PATCH] Use imports manager

---
 pkg/generate/generator.go     |  1 +
 pkg/translate/imports.go      | 31 +++++++++++++++++++++++++++++++
 pkg/translate/imports_test.go | 25 +++++++++++++++++++++++++
 templates/sdk/entry.tmpl      | 10 +---------
 templates/sdk/location.tmpl   |  8 +-------
 templates/sdk/service.tmpl    |  9 +--------
 6 files changed, 60 insertions(+), 24 deletions(-)
 create mode 100644 pkg/translate/imports.go
 create mode 100644 pkg/translate/imports_test.go

diff --git a/pkg/generate/generator.go b/pkg/generate/generator.go
index 99a2aba3..c623bb00 100644
--- a/pkg/generate/generator.go
+++ b/pkg/generate/generator.go
@@ -119,6 +119,7 @@ func (c *Creator) createFile(filePath string) (*os.File, error) {
 func (c *Creator) parseTemplate(templateName string) (*template.Template, error) {
 	templatePath := filepath.Join(c.TemplatesDir, templateName)
 	funcMap := template.FuncMap{
+		"renderImports":             translate.RenderImports,
 		"packageName":               translate.PackageName,
 		"locationType":              translate.LocationType,
 		"specParamType":             translate.SpecParamType,
diff --git a/pkg/translate/imports.go b/pkg/translate/imports.go
new file mode 100644
index 00000000..459195cf
--- /dev/null
+++ b/pkg/translate/imports.go
@@ -0,0 +1,31 @@
+package translate
+
+import "github.com/paloaltonetworks/pan-os-codegen/pkg/imports"
+
+func RenderImports(templateType string) (string, error) {
+	manager := imports.NewManager()
+
+	switch templateType {
+	case "entry":
+		manager.AddStandardImport("encoding/xml", "")
+		manager.AddStandardImport("fmt", "")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/filtering", "filtering")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/generic", "generic")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/util", "util")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/version", "version")
+	case "location":
+		manager.AddStandardImport("fmt", "")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/errors", "errors")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/util", "util")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/version", "version")
+	case "service":
+		manager.AddStandardImport("context", "")
+		manager.AddStandardImport("fmt", "")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/errors", "errors")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/filtering", "filtering")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/util", "util")
+		manager.AddSdkImport("github.com/PaloAltoNetworks/pango/xmlapi", "xmlapi")
+	}
+
+	return manager.RenderImports()
+}
diff --git a/pkg/translate/imports_test.go b/pkg/translate/imports_test.go
new file mode 100644
index 00000000..5c1456e4
--- /dev/null
+++ b/pkg/translate/imports_test.go
@@ -0,0 +1,25 @@
+package translate
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestRenderImports(t *testing.T) {
+	// given
+	expectedImports := `
+import (
+    "fmt"
+
+    errors "github.com/PaloAltoNetworks/pango/errors"
+    util "github.com/PaloAltoNetworks/pango/util"
+    version "github.com/PaloAltoNetworks/pango/version"
+)`
+
+	// when
+	actualImports, _ := RenderImports("location")
+
+	// then
+	assert.NotNil(t, actualImports)
+	assert.Equal(t, expectedImports, actualImports)
+}
diff --git a/templates/sdk/entry.tmpl b/templates/sdk/entry.tmpl
index acea588e..500ece69 100644
--- a/templates/sdk/entry.tmpl
+++ b/templates/sdk/entry.tmpl
@@ -1,15 +1,7 @@
 {{- if .Entry}}
     package {{packageName .GoSdkPath}}
 
-    import (
-    "encoding/xml"
-    "fmt"
-
-    "github.com/PaloAltoNetworks/pango/filtering"
-    "github.com/PaloAltoNetworks/pango/generic"
-    "github.com/PaloAltoNetworks/pango/util"
-    "github.com/PaloAltoNetworks/pango/version"
-    )
+    {{renderImports "entry"}}
 
     var (
     _ filtering.Fielder = &Entry{}
diff --git a/templates/sdk/location.tmpl b/templates/sdk/location.tmpl
index e23bebde..3323db24 100644
--- a/templates/sdk/location.tmpl
+++ b/templates/sdk/location.tmpl
@@ -1,12 +1,6 @@
 package {{packageName .GoSdkPath}}
 
-import (
-"fmt"
-
-"github.com/PaloAltoNetworks/pango/errors"
-"github.com/PaloAltoNetworks/pango/util"
-"github.com/PaloAltoNetworks/pango/version"
-)
+{{renderImports "location"}}
 
 type Location struct {
 {{range $key, $location := .Locations}}
diff --git a/templates/sdk/service.tmpl b/templates/sdk/service.tmpl
index a81358bc..36be98af 100644
--- a/templates/sdk/service.tmpl
+++ b/templates/sdk/service.tmpl
@@ -1,14 +1,7 @@
 package {{packageName .GoSdkPath}}
 {{- if .Entry}}
-    import (
-    "context"
-    "fmt"
+    {{renderImports "service"}}
 
-    "github.com/PaloAltoNetworks/pango/errors"
-    "github.com/PaloAltoNetworks/pango/filtering"
-    "github.com/PaloAltoNetworks/pango/util"
-    "github.com/PaloAltoNetworks/pango/xmlapi"
-    )
 
     type Service struct {
     client util.PangoClient