feat: support for decoding localstatequery queries (#813)

blinklabs-io · Dec 30, 2024 · 52b4697 · 52b4697
1 parent 70be436
commit 52b4697
Show file tree

Hide file tree

Showing 4 changed files with 338 additions and 11 deletions.
diff --git a/protocol/localstatequery/localstatequery.go b/protocol/localstatequery/localstatequery.go
@@ -155,7 +155,7 @@ type CallbackContext struct {
 
 // Callback function types
 type AcquireFunc func(CallbackContext, AcquireTarget, bool) error
-type QueryFunc func(CallbackContext, any) (any, error)
+type QueryFunc func(CallbackContext, QueryWrapper) (any, error)
 type ReleaseFunc func(CallbackContext) error
 
 // New returns a new LocalStateQuery object

diff --git a/protocol/localstatequery/messages.go b/protocol/localstatequery/messages.go
@@ -154,15 +154,17 @@ func NewMsgFailure(failure uint8) *MsgFailure {
 
 type MsgQuery struct {
 	protocol.MessageBase
-	Query interface{}
+	Query QueryWrapper
 }
 
 func NewMsgQuery(query interface{}) *MsgQuery {
 	m := &MsgQuery{
 		MessageBase: protocol.MessageBase{
 			MessageType: MessageTypeQuery,
 		},
-		Query: query,
+		Query: QueryWrapper{
+			Query: query,
+		},
 	}
 	return m
 }

diff --git a/protocol/localstatequery/messages_test.go b/protocol/localstatequery/messages_test.go
@@ -16,11 +16,12 @@ package localstatequery
 
 import (
 	"encoding/hex"
+	"reflect"
+	"testing"
+
 	"github.com/blinklabs-io/gouroboros/cbor"
 	"github.com/blinklabs-io/gouroboros/protocol"
 	"github.com/blinklabs-io/gouroboros/protocol/common"
-	"reflect"
-	"testing"
 )
 
 type testDefinition struct {
@@ -49,12 +50,12 @@ var tests = []testDefinition{
 		CborHex: "8203820082028101",
 		Message: NewMsgQuery(
 			// Current era hard-fork query
-			[]interface{}{
-				uint64(0),
-				[]interface{}{
-					uint64(2),
-					[]interface{}{
-						uint64(1),
+			&BlockQuery{
+				Query: &HardForkQuery{
+					Query: &HardForkCurrentEraQuery{
+						simpleQueryBase{
+							Type: QueryTypeHardForkCurrentEra,
+						},
 					},
 				},
 			},
@@ -105,6 +106,9 @@ func TestDecode(t *testing.T) {
 		}
 		// Set the raw CBOR so the comparison should succeed
 		test.Message.SetCbor(cborData)
+		if m, ok := msg.(*MsgQuery); ok {
+			m.Query.SetCbor(nil)
+		}
 		if !reflect.DeepEqual(msg, test.Message) {
 			t.Fatalf(
 				"CBOR did not decode to expected message object\n  got:    %#v\n  wanted: %#v",