diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index d735bc15..2e628b00 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -22,4 +22,6 @@ jobs:
         run: ls -alh build/bin
       - name: Test build
         run: echo w | ./build/bin/pokesay-linux-amd64
+      - name: Run unit tests
+        run: go test
 
diff --git a/build/Makefile b/build/Makefile
index ae3973d8..91f93404 100644
--- a/build/Makefile
+++ b/build/Makefile
@@ -24,18 +24,18 @@ build/cows:
 # generate embedded bin files for category/metadata/the actual pokemon
 build/assets:
 	docker run \
-		-v $(PWD)/build/assets:/usr/local/src/build/assets \
+		-v $(PWD)/assets:/usr/local/src/build/assets \
 		--rm --name pokesay \
 	    $(DOCKER_IMAGE) \
 			build/scripts/build_assets.sh
 
 build/release: build/assets
 	docker run \
-		-v $(PWD)/build/bin:/usr/local/src/build/bin \
-		-v $(PWD)/build/assets:/usr/local/src/build/assets \
+		-v $(PWD)/bin:/usr/local/src/build/bin \
+		-v $(PWD)/assets:/usr/local/src/build/assets \
 		--rm --name pokesay \
 		  $(DOCKER_IMAGE) \
 		  build/scripts/build.sh
-	tree $(PWD)/build/bin/
+	tree $(PWD)/bin/
 
 .PHONY: all build/docker build/cows build/assets build/release
diff --git a/build/scripts/build_assets.sh b/build/scripts/build_assets.sh
index 0c175adb..e5f73ba1 100755
--- a/build/scripts/build_assets.sh
+++ b/build/scripts/build_assets.sh
@@ -2,9 +2,10 @@
 
 set -euo pipefail
 
-tar xzf build/cows.tar.gz
+FROM="${1:-/tmp/cows/}"
+
 go run ./src/bin/pokedex/pokedex.go \
-  -from ./cows/ \
+  -from "${FROM}" \
   -to ./build/assets/ \
   -toCategoryFpath pokedex.gob \
   -toDataSubDir cows/ \
diff --git a/pokesay.go b/pokesay.go
index b5a09f0a..a1b7239b 100644
--- a/pokesay.go
+++ b/pokesay.go
@@ -3,11 +3,13 @@ package main
 import (
 	"bufio"
 	"embed"
+	"errors"
 	"flag"
 	"fmt"
 	"log"
 	"math/rand"
 	"os"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -41,6 +43,9 @@ func check(e error) {
 }
 
 func randomInt(n int) int {
+	if n <= 0 {
+		log.Fatal(errors.New("randomInt arg must be >0"))
+	}
 	return rand.New(Rand).Intn(n)
 }
 
@@ -145,7 +150,7 @@ func parseFlags() Args {
 	return args
 }
 
-func runListCategories(categories pokedex.PokemonTrie) {
+func ListCategories(categories pokedex.PokemonTrie) []string {
 	ukm := map[string]bool{}
 	for _, v := range categories.Keys {
 		for _, k := range v {
@@ -158,8 +163,13 @@ func runListCategories(categories pokedex.PokemonTrie) {
 		keys[i] = k
 		i++
 	}
-	fmt.Println(strings.Join(keys, " "))
-	fmt.Printf("\n%d %s\n", len(keys), "total names")
+	sort.Strings(keys)
+	return keys
+}
+
+func runListCategories(categories pokedex.PokemonTrie) {
+	keys := ListCategories(categories)
+	fmt.Printf("%s\n%d %s\n", strings.Join(keys, " "), len(keys), "total categories")
 }
 
 func runListNames() {
diff --git a/pokesay_test.go b/pokesay_test.go
new file mode 100644
index 00000000..7bd955a9
--- /dev/null
+++ b/pokesay_test.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/tmck-code/pokesay/src/pokedex"
+)
+
+// Made my own basic test helper. Takes in an expected & result object of any type, and asserts
+// that their Go syntax representations (%#v) are the same
+func assert(expected interface{}, result interface{}, obj interface{}, test *testing.T) {
+	if fmt.Sprintf("%#v", expected) != fmt.Sprintf("%#v", result) {
+		test.Fatalf(`expected = %#v, result = %#v, obj = %#v`, expected, result, obj)
+	}
+}
+
+func TestNewPokemonEntry(test *testing.T) {
+	p := pokedex.NewPokemonEntry(1, "yo")
+	assert(1, p.Index, p, test)
+}
+
+func TestTrieInsert(test *testing.T) {
+	t := pokedex.NewTrie()
+	t.Insert([]string{"p", "g1", "r"}, pokedex.NewPokemonEntry(0, "pikachu"))
+	t.Insert([]string{"p", "g1", "r"}, pokedex.NewPokemonEntry(1, "bulbasaur"))
+
+	result, err := t.GetCategory([]string{"p", "g1"})
+	check(err)
+
+	assert(2, len(result), result, test)
+	assert(
+		"[{Index: 0, Name: pikachu} {Index: 1, Name: bulbasaur}]",
+		fmt.Sprintf("%s", result),
+		result, test,
+	)
+}
+
+func TestCategoryPaths(test *testing.T) {
+	t := pokedex.NewTrie()
+	t.Insert([]string{"small", "g1", "r"}, pokedex.NewPokemonEntry(0, "pikachu"))
+	t.Insert([]string{"small", "g1", "o"}, pokedex.NewPokemonEntry(1, "bulbasaur"))
+	t.Insert([]string{"medium", "g1", "o"}, pokedex.NewPokemonEntry(2, "bulbasaur"))
+	t.Insert([]string{"big", "g1", "o"}, pokedex.NewPokemonEntry(3, "bulbasaur"))
+	t.Insert([]string{"big", "g1"}, pokedex.NewPokemonEntry(4, "charmander"))
+
+	expected := [][]string{
+		[]string{"small", "g1", "r"},
+		[]string{"small", "g1", "o"},
+		[]string{"medium", "g1", "o"},
+		[]string{"big", "g1", "o"},
+		[]string{"big", "g1"},
+	}
+	assert(expected, t.Keys, t, test)
+
+	expected = [][]string{
+		[]string{"big", "g1", "o"},
+		[]string{"big", "g1"},
+	}
+	result, err := t.GetCategoryPaths("big")
+	check(err)
+	assert(expected, result, result, test)
+}
+
+func TestListCategories(test *testing.T) {
+	t := pokedex.NewTrie()
+	t.Insert([]string{"small", "g1", "r"}, pokedex.NewPokemonEntry(0, "pikachu"))
+	t.Insert([]string{"small", "g1", "o"}, pokedex.NewPokemonEntry(1, "bulbasaur"))
+	t.Insert([]string{"medium", "g1", "o"}, pokedex.NewPokemonEntry(2, "bulbasaur"))
+	t.Insert([]string{"big", "g1", "o"}, pokedex.NewPokemonEntry(3, "bulbasaur"))
+	t.Insert([]string{"big", "g1"}, pokedex.NewPokemonEntry(4, "charmander"))
+
+	result := ListCategories(*t)
+	assert([]string{"big", "g1", "medium", "o", "r", "small"}, result, result, test)
+}
diff --git a/src/pokedex/convert.go b/src/pokedex/convert.go
index ed4179fd..28bca714 100644
--- a/src/pokedex/convert.go
+++ b/src/pokedex/convert.go
@@ -130,7 +130,7 @@ func CreateMetadata(rootDir string, fpaths []string, debug bool) (PokemonTrie, [
 		data, err := os.ReadFile(fpath)
 		check(err)
 
-		cats := createCategories(strings.TrimPrefix(fpath, rootDir))
+		cats := createCategories(strings.TrimPrefix(fpath, rootDir), data)
 		name := createName(fpath)
 
 		categories.Insert(
@@ -147,9 +147,20 @@ func createName(fpath string) string {
 	return strings.Split(parts[len(parts)-1], ".")[0]
 }
 
-func createCategories(fpath string) []string {
+func SizeCategory(height int) string {
+	if height <= 13 {
+		return "small"
+	} else if height <= 19 {
+		return "medium"
+	}
+	return "big"
+}
+
+func createCategories(fpath string, data []byte) []string {
 	parts := strings.Split(fpath, "/")
-	return parts[0 : len(parts)-1]
+	height := SizeCategory(len(strings.Split(string(data), "\n")))
+
+	return append([]string{height}, parts[0:len(parts)-1]...)
 }
 
 // Strips the leading "./" from a path e.g. "./cows/ -> cows/"
diff --git a/src/pokedex/entries.go b/src/pokedex/entries.go
index a1c94fd7..1473143f 100644
--- a/src/pokedex/entries.go
+++ b/src/pokedex/entries.go
@@ -24,6 +24,11 @@ type PokemonEntry struct {
 	Index int
 }
 
+func (p PokemonEntry) String() string {
+	return fmt.Sprintf("{Index: %d, Name: %s}", p.Index, p.Name)
+}
+
+
 type Node struct {
 	Children map[string]*Node
 	Data     []*PokemonEntry
@@ -109,13 +114,9 @@ func (t *PokemonTrie) Insert(s []string, data *PokemonEntry) {
 func (t PokemonTrie) GetCategoryPaths(s string) ([][]string, error) {
 	matches := [][]string{}
 	for _, k := range t.Keys {
-		for i, el := range k {
+		for _, el := range k {
 			if el == s {
-				if i == 0 {
-					matches = append(matches, []string{s})
-				} else {
-					matches = append(matches, k)
-				}
+				matches = append(matches, k)
 			}
 		}
 	}
@@ -138,6 +139,9 @@ func (t PokemonTrie) GetCategory(s []string) ([]*PokemonEntry, error) {
 			return nil, errors.New(fmt.Sprintf("Could not find category: %s", s))
 		}
 	}
+	if len(matches) == 0 {
+		return nil, errors.New(fmt.Sprintf("Could not find category: %s", s))
+	}
 	return matches, nil
 }