From f15065e1545211032719c95ecf92d4648b6d30bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20F=2E=20R=C3=B8dseth?= <rodseth@gmail.com>
Date: Tue, 14 May 2024 11:19:55 +0200
Subject: [PATCH] Add tests

---
 convert_test.go  | 57 ++++++++++++++++++++++++++++++++++++++++++++++--
 generate_test.go | 30 +++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/convert_test.go b/convert_test.go
index 9a0755e..e9a4be3 100644
--- a/convert_test.go
+++ b/convert_test.go
@@ -49,7 +49,7 @@ func TestConvert(t *testing.T) {
 }
 
 func TestConvertFewColors(t *testing.T) {
-	for _, N := range []int{8, 16, 32, 64, 128} {
+	for _, N := range []int{2, 4, 8, 16, 32, 64, 128} {
 
 		// Read a True Color PNG file
 		data, err := os.Open("testdata/splash.png")
@@ -142,7 +142,7 @@ func TestConvertGeneral(t *testing.T) {
 }
 
 func TestConvertMountain(t *testing.T) {
-	for _, N := range []int{8, 16, 32, 64, 128, 256} {
+	for _, N := range []int{2, 4, 8, 16, 32, 64, 128, 256} {
 
 		// Read a True Color PNG file
 		data, err := os.Open("testdata/tm_small.png")
@@ -194,6 +194,59 @@ func TestConvertMountain(t *testing.T) {
 	}
 }
 
+func TestConvertRainforest(t *testing.T) {
+	for _, N := range []int{2, 4, 8, 16, 32, 64, 128, 256} {
+
+		// Read a True Color PNG file
+		data, err := os.Open("testdata/rainforest.png")
+		if err != nil {
+			t.Error(err)
+		}
+
+		// Decode the PNG image
+		img, err := png.Decode(data)
+		if err != nil {
+			t.Error(err)
+		}
+
+		// Generate a palette with N colors
+		pal, err := Generate(img, N)
+		if err != nil {
+			t.Error(err)
+		}
+
+		if len(pal) != N {
+			t.Fatalf("The palette should be %d long, but it is %d long.\n", N, len(pal))
+		}
+
+		// Convert the image to only use the given palette
+		imgN, err := ConvertCustom(img, pal)
+		if err != nil {
+			t.Error(err)
+		}
+
+		_, ok := imgN.(image.PalettedImage)
+		if !ok {
+			t.Fatal("The image should be an image.PalettedImage")
+		}
+
+		// Output the indexed image
+		f, err := os.Create("testdata/rainforest" + strconv.Itoa(N) + ".png")
+		if err != nil {
+			t.Error(err)
+		}
+
+		if err := png.Encode(f, imgN); err != nil {
+			f.Close()
+			t.Error(err)
+		}
+
+		if err := f.Close(); err != nil {
+			t.Error(err)
+		}
+	}
+}
+
 func TestConvertPlan9(t *testing.T) {
 	// Read a True Color PNG file
 	data, err := os.Open("testdata/splash.png")
diff --git a/generate_test.go b/generate_test.go
index 985938f..ab67a4e 100644
--- a/generate_test.go
+++ b/generate_test.go
@@ -1,6 +1,7 @@
 package palgen
 
 import (
+	"fmt"
 	"image/png"
 	"os"
 	"testing"
@@ -104,3 +105,32 @@ func TestLarge(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestSmallImagePalette(t *testing.T) {
+	// Read images and try to generate palettes of exactly 2 and 4 colors
+	for _, n := range []int{2, 4} {
+		for _, imageName := range []string{"rainforest", "splash", "tm"} {
+			data, err := os.Open(fmt.Sprintf("testdata/%s.png", imageName))
+			if err != nil {
+				t.Errorf("Failed to open image file: %v", err)
+			}
+			defer data.Close()
+
+			img, err := png.Decode(data)
+			if err != nil {
+				t.Errorf("Failed to decode PNG image: %v", err)
+			}
+
+			// Generate a palette of N colors
+			pal, err := Generate(img, n)
+			if err != nil {
+				t.Errorf("Failed to generate a palette of %d colors: %v", n, err)
+			}
+
+			// Check if the generated palette has exactly N colors
+			if len(pal) != n {
+				t.Errorf("Generated palette for testdata/%s.png has %d colors, expected %d", imageName, len(pal), n)
+			}
+		}
+	}
+}