Add canister id option to generator.

Author: q-uint

cmd/goic/main.go

@@ -71,6 +71,14 @@ var root = cmd.NewCommandFork(
 					Name:     "packageName",
 					HasValue: true,
 				},
+				{
+					Name:     "agentName",
+					HasValue: true,
+				},
+				{
+					Name:     "canisterID",
+					HasValue: true,
+				},
 				{
 					Name:     "indirect",
 					HasValue: false,
@@ -94,8 +102,22 @@ var root = cmd.NewCommandFork(
 					packageName = p
 				}
 
+				var agentName string
+				if a, ok := options["agentName"]; ok {
+					agentName = a
+				}
+
+				var canisterID *principal.Principal
+				if cID, ok := options["canisterID"]; ok {
+					p, err := principal.Decode(cID)
+					if err != nil {
+						return err
+					}
+					canisterID = &p
+				}
+
 				_, indirect := options["indirect"]
-				return writeDID(canisterName, packageName, path, []rune(string(rawDID)), indirect)
+				return writeDID(agentName, canisterName, canisterID, packageName, path, []rune(string(rawDID)), indirect)
 			},
 		),
 		cmd.NewCommand(
@@ -111,18 +133,22 @@ var root = cmd.NewCommandFork(
 					Name:     "packageName",
 					HasValue: true,
 				},
+				{
+					Name:     "agentName",
+					HasValue: true,
+				},
 				{
 					Name:     "indirect",
 					HasValue: false,
 				},
 			},
 			func(args []string, options map[string]string) error {
 				id := args[0]
-				canisterId, err := principal.Decode(id)
+				canisterID, err := principal.Decode(id)
 				if err != nil {
 					return err
 				}
-				rawDID, err := fetchDID(canisterId)
+				rawDID, err := fetchDID(canisterID)
 				if err != nil {
 					return err
 				}
@@ -138,8 +164,13 @@ var root = cmd.NewCommandFork(
 					packageName = p
 				}
 
+				var agentName string
+				if a, ok := options["agentName"]; ok {
+					agentName = a
+				}
+
 				_, indirect := options["indirect"]
-				return writeDID(canisterName, packageName, path, []rune(string(rawDID)), indirect)
+				return writeDID(agentName, canisterName, &canisterID, packageName, path, []rune(string(rawDID)), indirect)
 			},
 		),
 	),
@@ -167,14 +198,17 @@ func main() {
 	}
 }
 
-func writeDID(canisterName, packageName, outputPath string, rawDID []rune, indirect bool) error {
-	g, err := gen.NewGenerator("", canisterName, packageName, rawDID)
+func writeDID(agentName, canisterName string, canisterID *principal.Principal, packageName, outputPath string, rawDID []rune, indirect bool) error {
+	g, err := gen.NewGenerator(agentName, canisterName, packageName, rawDID)
 	if err != nil {
 		return err
 	}
 	if indirect {
 		g.Indirect()
 	}
+	if canisterID != nil {
+		g.WithCanisterID(canisterID)
+	}
 	raw, err := g.Generate()
 	if err != nil {
 		return err

gen/generator.go

@@ -11,6 +11,7 @@ import (
 	"unicode"
 
 	"github.com/aviate-labs/agent-go/candid/did"
+	"github.com/aviate-labs/agent-go/principal"
 )
 
 const (
@@ -72,6 +73,7 @@ type Generator struct {
 	AgentName          string
 	ModulePath         string
 	CanisterName       string
+	CanisterID         *principal.Principal
 	PackageName        string
 	ServiceDescription did.Description
 	usedIDL            bool
@@ -159,12 +161,14 @@ func (g *Generator) Generate() ([]byte, error) {
 	}
 	var tmpl bytes.Buffer
 	if err := t.Execute(&tmpl, agentArgs{
-		AgentName:    g.AgentName,
-		CanisterName: g.CanisterName,
-		PackageName:  g.PackageName,
-		UsedIDL:      g.usedIDL,
-		Definitions:  definitions,
-		Methods:      methods,
+		AgentName:      g.AgentName,
+		AgentNameUpper: strings.ToUpper(g.AgentName),
+		CanisterName:   g.CanisterName,
+		CanisterID:     g.CanisterID,
+		PackageName:    g.PackageName,
+		UsedIDL:        g.usedIDL,
+		Definitions:    definitions,
+		Methods:        methods,
 	}); err != nil {
 		return nil, err
 	}
@@ -177,6 +181,11 @@ func (g *Generator) Indirect() *Generator {
 	return g
 }
 
+func (g *Generator) WithCanisterID(canisterID *principal.Principal) *Generator {
+	g.CanisterID = canisterID
+	return g
+}
+
 func (g *Generator) dataToString(prefix string, data did.Data) string {
 	switch t := data.(type) {
 	case did.Blob:
@@ -325,12 +334,14 @@ func (g *Generator) dataToString(prefix string, data did.Data) string {
 }
 
 type agentArgs struct {
-	AgentName    string
-	CanisterName string
-	PackageName  string
-	UsedIDL      bool
-	Definitions  []agentArgsDefinition
-	Methods      []agentArgsMethod
+	AgentName      string
+	AgentNameUpper string
+	CanisterName   string
+	CanisterID     *principal.Principal
+	PackageName    string
+	UsedIDL        bool
+	Definitions    []agentArgsDefinition
+	Methods        []agentArgsMethod
 }
 
 type agentArgsDefinition struct {

gen/templates/agent.gotmpl

@@ -7,6 +7,9 @@ import (
     "github.com/aviate-labs/agent-go/candid/idl"{{ end }}
     "github.com/aviate-labs/agent-go/principal"
 )
+{{- if .CanisterID }}
+
+var {{ if .AgentNameUpper }}{{ .AgentNameUpper }}_{{ end }}CANISTER_ID = principal.MustDecode("{{ .CanisterID.String }}"){{ end }}
 
 {{- range .Definitions }}
 

gen/templates/agent_indirect.gotmpl

@@ -7,6 +7,9 @@ import (
     "github.com/aviate-labs/agent-go/candid/idl"{{ end }}
     "github.com/aviate-labs/agent-go/principal"
 )
+{{- if .CanisterID }}
+
+var {{ if .AgentNameUpper }}{{ .AgentNameUpper }}_{{ end }}CANISTER_ID = principal.MustDecode("{{ .CanisterID.String }}"){{ end }}
 
 {{- range .Definitions }}
 

subnet_test.go

@@ -1,9 +1,10 @@
 package agent
 
 import (
+	"testing"
+
 	"github.com/aviate-labs/agent-go/certification"
 	"github.com/aviate-labs/agent-go/principal"
-	"testing"
 )
 
 func TestAgent_GetSubnetMetrics(t *testing.T) {