Add subnet list test.

aviate-labs · May 14, 2024 · 51e4c9b · 51e4c9b
1 parent 83f34a7
commit 51e4c9b
Show file tree

Hide file tree

Showing 7 changed files with 180 additions and 141 deletions.
diff --git a/agent.go b/agent.go
@@ -233,15 +233,7 @@ func (a Agent) GetCanisterControllers(canisterID principal.Principal) ([]princip
 // GetCanisterInfo returns the raw certificate for the given canister based on the given sub-path.
 func (a Agent) GetCanisterInfo(canisterID principal.Principal, subPath string) ([]byte, error) {
 	path := []hashtree.Label{hashtree.Label("canister"), canisterID.Raw, hashtree.Label(subPath)}
-	c, err := a.readStateCertificate(canisterID, [][]hashtree.Label{path})
-	if err != nil {
-		return nil, err
-	}
-	var state map[string]any
-	if err := cbor.Unmarshal(c, &state); err != nil {
-		return nil, err
-	}
-	node, err := hashtree.DeserializeNode(state["tree"].([]any))
+	node, err := a.ReadState(canisterID, [][]hashtree.Label{path})
 	if err != nil {
 		return nil, err
 	}
@@ -298,6 +290,19 @@ func (a Agent) Query(canisterID principal.Principal, methodName string, args []a
 	return query.Query(values...)
 }
 
+// ReadState reads the state of the given canister at the given path.
+func (a Agent) ReadState(canisterID principal.Principal, path [][]hashtree.Label) (hashtree.Node, error) {
+	c, err := a.readStateCertificate(canisterID, path)
+	if err != nil {
+		return nil, err
+	}
+	var state map[string]any
+	if err := cbor.Unmarshal(c, &state); err != nil {
+		return nil, err
+	}
+	return hashtree.DeserializeNode(state["tree"].([]any))
+}
+
 // RequestStatus returns the status of the request with the given ID.
 func (a Agent) RequestStatus(ecID principal.Principal, requestID RequestID) ([]byte, hashtree.Node, error) {
 	a.logger.Printf("[AGENT] REQUEST STATUS %s %x", ecID, requestID)

diff --git a/agent_test.go b/agent_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/aviate-labs/agent-go"
 	"github.com/aviate-labs/agent-go/candid/idl"
+	"github.com/aviate-labs/agent-go/certification/hashtree"
 	"github.com/aviate-labs/agent-go/ic"
 	mgmt "github.com/aviate-labs/agent-go/ic/ic"
 	"github.com/aviate-labs/agent-go/ic/icpledger"
@@ -98,6 +99,23 @@ func Example_query_secp256k1() {
 	// 0
 }
 
+func TestAgent_Call(t *testing.T) {
+	a, err := agent.New(agent.DefaultConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+	n, err := a.ReadState(ic.REGISTRY_PRINCIPAL, [][]hashtree.Label{{hashtree.Label("subnet")}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, path := range hashtree.ListPaths(n, nil) {
+		if len(path) == 3 && string(path[0]) == "subnet" && string(path[2]) == "public_key" {
+			subnetID := principal.Principal{Raw: []byte(path[1])}
+			_ = subnetID
+		}
+	}
+}
+
 func TestAgent_Call_bitcoinGetBalanceQuery(t *testing.T) {
 	a, err := mgmt.NewAgent(ic.MANAGEMENT_CANISTER_PRINCIPAL, agent.DefaultConfig)
 	if err != nil {

diff --git a/certification/hashtree/node.go b/certification/hashtree/node.go
@@ -15,6 +15,22 @@ func DomainSeparator(t string) []byte {
 	)
 }
 
+// ListPaths returns all paths from the root to a leaf node.
+func ListPaths(n Node, path []Label) [][]Label {
+	switch n := n.(type) {
+	case Fork:
+		l := ListPaths(n.LeftTree, path)
+		r := ListPaths(n.RightTree, path)
+		return append(l, r...)
+	case Labeled:
+		return ListPaths(n.Tree, append(path, n.Label))
+	case Leaf:
+		return [][]Label{path}
+	default: // Empty, Pruned
+		return nil
+	}
+}
+
 func Serialize(node Node) ([]byte, error) {
 	return cbor.Marshal(serialize(node))
 }

diff --git a/ic/sns/ledger/agent.go b/ic/sns/ledger/agent.go
@@ -485,7 +485,7 @@ type GetBlocksResult struct {
 	} `ic:"blocks" json:"blocks"`
 	ArchivedBlocks []struct {
 		Args     []GetBlocksArgs `ic:"args" json:"args"`
-		Callback struct { /* NOT SUPPORTED */
+		Callback struct {        /* NOT SUPPORTED */
 		} `ic:"callback" json:"callback"`
 	} `ic:"archived_blocks" json:"archived_blocks"`
 }

diff --git a/pocketic/pocketic.go b/pocketic/pocketic.go
@@ -419,11 +419,11 @@ type Topology struct {
 	CanisterRanges []CanisterIDRange `json:"canister_ranges"`
 }
 
-type rawCanisterID struct {
-	CanisterID string `json:"canister_id"`
-}
-
 type canisterIDRange struct {
 	Start rawCanisterID `json:"start"`
 	End   rawCanisterID `json:"end"`
 }
+
+type rawCanisterID struct {
+	CanisterID string `json:"canister_id"`
+}
diff --git a/pocketic/raw.go b/pocketic/raw.go
@@ -38,57 +38,6 @@ func (b *Base64EncodedBlob) UnmarshalJSON(bytes []byte) error {
 	return nil
 }
 
-type createResponse[T any] struct {
-	Created *T            `json:"Created"`
-	Error   *ErrorMessage `json:"Error"`
-}
-
-type rawAddCycles struct {
-	Amount     int               `json:"amount"`
-	CanisterID Base64EncodedBlob `json:"canister_id"`
-}
-
-type rawCanisterCall struct {
-	CanisterID         Base64EncodedBlob  `json:"canister_id"`
-	EffectivePrincipal EffectivePrincipal `json:"effective_principal"`
-	Method             string             `json:"method"`
-	Payload            Base64EncodedBlob  `json:"payload"`
-	Sender             Base64EncodedBlob  `json:"sender"`
-}
-
-func (r *rawCanisterCall) UnmarshalJSON(bytes []byte) error {
-	var raw struct {
-		CanisterID         Base64EncodedBlob `json:"canister_id"`
-		EffectivePrincipal json.RawMessage   `json:"effective_principal"`
-		Method             string            `json:"method"`
-		Payload            Base64EncodedBlob `json:"payload"`
-		Sender             Base64EncodedBlob `json:"sender"`
-	}
-	if err := json.Unmarshal(bytes, &raw); err != nil {
-		return err
-	}
-	ep, err := unmarshalRawEffectivePrincipal(raw.EffectivePrincipal)
-	if err != nil {
-		return err
-	}
-	r.CanisterID = raw.CanisterID
-	r.EffectivePrincipal = ep
-	r.Method = raw.Method
-	r.Payload = raw.Payload
-	r.Sender = raw.Sender
-	return nil
-}
-
-type RawCanisterID struct {
-	CanisterID Base64EncodedBlob `json:"canister_id"`
-}
-
-type rawCanisterResult result[rawWasmResult]
-
-type rawCycles struct {
-	Cycles int `json:"cycles"`
-}
-
 type EffectivePrincipal interface {
 	rawEffectivePrincipal()
 }
@@ -158,23 +107,27 @@ func (r *MessageID) UnmarshalJSON(bytes []byte) error {
 	return nil
 }
 
-type rawSetStableMemory struct {
-	BlobID     Base64EncodedBlob `json:"blob_id"`
+type RawCanisterID struct {
 	CanisterID Base64EncodedBlob `json:"canister_id"`
 }
 
-type rawStableMemory struct {
-	Blob Base64EncodedBlob `json:"blob"`
-}
+type Reject string
 
-type rawSubmitIngressResult result[MessageID]
+func (r Reject) Error() string {
+	return string(r)
+}
 
 type SubnetID struct {
 	SubnetID Base64EncodedBlob `json:"subnet_id"`
 }
 
-type rawTime struct {
-	NanosSinceEpoch int64 `json:"nanos_since_epoch"`
+type UserError struct {
+	Code        string `json:"code"`
+	Description string `json:"description"`
+}
+
+func (e UserError) Error() string {
+	return fmt.Sprintf("(%s) %s", e.Code, e.Description)
 }
 
 type VerifyCanisterSigArg struct {
@@ -184,26 +137,73 @@ type VerifyCanisterSigArg struct {
 	Signature  Base64EncodedBlob `json:"sig"`
 }
 
-type rawWasmResult wasmResult[Base64EncodedBlob]
+type createResponse[T any] struct {
+	Created *T            `json:"Created"`
+	Error   *ErrorMessage `json:"Error"`
+}
 
-type Reject string
+type rawAddCycles struct {
+	Amount     int               `json:"amount"`
+	CanisterID Base64EncodedBlob `json:"canister_id"`
+}
 
-func (r Reject) Error() string {
-	return string(r)
+type rawCanisterCall struct {
+	CanisterID         Base64EncodedBlob  `json:"canister_id"`
+	EffectivePrincipal EffectivePrincipal `json:"effective_principal"`
+	Method             string             `json:"method"`
+	Payload            Base64EncodedBlob  `json:"payload"`
+	Sender             Base64EncodedBlob  `json:"sender"`
 }
 
-type result[R any] struct {
-	Ok  *R         `json:"Ok,omitempty"`
-	Err *UserError `json:"Err,omitempty"`
+func (r *rawCanisterCall) UnmarshalJSON(bytes []byte) error {
+	var raw struct {
+		CanisterID         Base64EncodedBlob `json:"canister_id"`
+		EffectivePrincipal json.RawMessage   `json:"effective_principal"`
+		Method             string            `json:"method"`
+		Payload            Base64EncodedBlob `json:"payload"`
+		Sender             Base64EncodedBlob `json:"sender"`
+	}
+	if err := json.Unmarshal(bytes, &raw); err != nil {
+		return err
+	}
+	ep, err := unmarshalRawEffectivePrincipal(raw.EffectivePrincipal)
+	if err != nil {
+		return err
+	}
+	r.CanisterID = raw.CanisterID
+	r.EffectivePrincipal = ep
+	r.Method = raw.Method
+	r.Payload = raw.Payload
+	r.Sender = raw.Sender
+	return nil
 }
 
-type UserError struct {
-	Code        string `json:"code"`
-	Description string `json:"description"`
+type rawCanisterResult result[rawWasmResult]
+
+type rawCycles struct {
+	Cycles int `json:"cycles"`
 }
 
-func (e UserError) Error() string {
-	return fmt.Sprintf("(%s) %s", e.Code, e.Description)
+type rawSetStableMemory struct {
+	BlobID     Base64EncodedBlob `json:"blob_id"`
+	CanisterID Base64EncodedBlob `json:"canister_id"`
+}
+
+type rawStableMemory struct {
+	Blob Base64EncodedBlob `json:"blob"`
+}
+
+type rawSubmitIngressResult result[MessageID]
+
+type rawTime struct {
+	NanosSinceEpoch int64 `json:"nanos_since_epoch"`
+}
+
+type rawWasmResult wasmResult[Base64EncodedBlob]
+
+type result[R any] struct {
+	Ok  *R         `json:"Ok,omitempty"`
+	Err *UserError `json:"Err,omitempty"`
 }
 
 // wasmResult describes the different types that executing a WASM function in a canister can produce.