From 5f9b611921c82eed94d33a1bfbc71f17f755c989 Mon Sep 17 00:00:00 2001
From: DevMiner <devminer@devminer.xyz>
Date: Wed, 18 Sep 2024 23:03:21 +0200
Subject: [PATCH] refactor: parse extension keys into proper array

---
 .../{lysand_entity.go => versia_entity.go}    |  0
 pkg/versia/entity.go                          |  5 ---
 pkg/versia/extension.go                       | 41 +++++++++++++++++++
 pkg/versia/extension_test.go                  | 29 +++++++++++++
 4 files changed, 70 insertions(+), 5 deletions(-)
 rename ent/schema/{lysand_entity.go => versia_entity.go} (100%)
 create mode 100644 pkg/versia/extension.go
 create mode 100644 pkg/versia/extension_test.go

diff --git a/ent/schema/lysand_entity.go b/ent/schema/versia_entity.go
similarity index 100%
rename from ent/schema/lysand_entity.go
rename to ent/schema/versia_entity.go
diff --git a/pkg/versia/entity.go b/pkg/versia/entity.go
index 702815a..8e74a96 100644
--- a/pkg/versia/entity.go
+++ b/pkg/versia/entity.go
@@ -24,8 +24,3 @@ type Entity struct {
 	// Extensions is a map of active extensions for the entity
 	Extensions Extensions `json:"extensions,omitempty"`
 }
-
-// Extensions represents the active extensions on an entity. For more information, see the [Spec].
-//
-// [Spec]: https://versia.pub/extensions#extension-definition
-type Extensions map[string]any
diff --git a/pkg/versia/extension.go b/pkg/versia/extension.go
new file mode 100644
index 0000000..bad0869
--- /dev/null
+++ b/pkg/versia/extension.go
@@ -0,0 +1,41 @@
+package versia
+
+import (
+	"fmt"
+	"strings"
+)
+
+// Extensions represents the active extensions on an entity. For more information, see the [Spec].
+//
+// [Spec]: https://versia.pub/extensions#extension-definition
+type Extensions map[ExtensionKey]any
+
+// ExtensionKey represents an extension's key. For more information see the [Spec].
+//
+// [Spec]: https://versia.pub/types#extensions
+type ExtensionKey [2]string
+
+func (e *ExtensionKey) UnmarshalText(b []byte) (err error) {
+	raw := string(b)
+
+	spl := strings.Split(raw, ":")
+	if len(spl) != 2 {
+		return InvalidExtensionKeyError{raw}
+	}
+
+	*e = [2]string{spl[0], spl[1]}
+
+	return
+}
+
+func (e ExtensionKey) MarshalText() ([]byte, error) {
+	return []byte(fmt.Sprintf("%s:%s", e[0], e[1])), nil
+}
+
+type InvalidExtensionKeyError struct {
+	Raw string
+}
+
+func (e InvalidExtensionKeyError) Error() string {
+	return fmt.Sprintf("invalid extension key: %s", e.Raw)
+}
diff --git a/pkg/versia/extension_test.go b/pkg/versia/extension_test.go
new file mode 100644
index 0000000..109945b
--- /dev/null
+++ b/pkg/versia/extension_test.go
@@ -0,0 +1,29 @@
+package versia
+
+import (
+	"encoding/json"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestExtensionKey_UnmarshalJSON(t *testing.T) {
+	cases := []struct {
+		Raw      string
+		Expected ExtensionKey
+		Error    bool
+	}{
+		{"\"pub.versia:Emoji\"", ExtensionKey{"pub.versia", "Emoji"}, false},
+		{"\"pub.versia\"", ExtensionKey{}, true},
+	}
+
+	for _, case_ := range cases {
+		key := ExtensionKey{}
+		err := json.Unmarshal([]byte(case_.Raw), &key)
+
+		assert.Equal(t, case_.Error, err != nil, "error assertion should match")
+
+		if !case_.Error {
+			assert.Equal(t, case_.Expected, key)
+		}
+	}
+}