Add tests for Prime Factorization function

2 files changed, 50 insertions, 0 deletions

diff --git a/factors/factors_test.go b/factors/factors_test.go
new file mode 100644
index 0000000..472a2e3
--- /dev/null
+++ b/factors/factors_test.go
@@ -0,0 +1,46 @@
+package factors
+
+import (
+	"fmt"
+	"testing"
+)
+
+var primeFactorCases = map[uint]PrimeFactorization{
+	0: PrimeFactorization{map[uint]uint{0: 1}},
+	1: PrimeFactorization{map[uint]uint{}},
+	2: PrimeFactorization{map[uint]uint{2: 1}},
+	3: PrimeFactorization{map[uint]uint{3: 1}},
+	4: PrimeFactorization{map[uint]uint{2: 2}},
+	6: PrimeFactorization{map[uint]uint{2: 1, 3: 1}},
+	10: PrimeFactorization{map[uint]uint{2: 1, 5: 1}},
+	12: PrimeFactorization{map[uint]uint{2: 2, 3: 1}},
+	33: PrimeFactorization{map[uint]uint{3: 1, 11: 1}},
+	60: PrimeFactorization{map[uint]uint{2: 2, 3: 1, 5: 1}},
+	86400: PrimeFactorization{map[uint]uint{2: 7, 3: 3, 5: 2}},
+}
+
+func TestPrimeFactorize(t *testing.T) {
+	for i, expected := range primeFactorCases {
+		testname := fmt.Sprintf("%d", i)
+		t.Run(testname, func(t *testing.T) {
+			actual := PrimeFactorize(i)
+			if !mapEquals(expected.exponents, actual.exponents) {
+				t.Errorf("PrimeFactorize(%d) = %s, want %s", i, actual, expected)
+			}
+		})
+	}
+}
+
+func mapEquals(a, b map[uint]uint) bool {
+	for k := range a {
+		if a[k] != b[k] {
+			return false
+		}
+	}
+	for k := range b {
+		if a[k] != b[k] {
+			return false
+		}
+	}
+	return true
+}
diff --git a/factors/prime_factorization.go b/factors/prime_factorization.go
index 62498b5..3f3abc7 100644
--- a/factors/prime_factorization.go
+++ b/factors/prime_factorization.go
@@ -15,6 +15,10 @@ type PrimeFactorization struct {
 
 // PrimeFactorize creates a PrimeFactorization by factorizing a number.
 func PrimeFactorize(n uint) PrimeFactorization {
+	if n == 0 {
+		return PrimeFactorization{map[uint]uint{0: 1}}
+	}
+
 	unfactored := n // number with all found factors removed (divided)
 	exponents := make(map[uint]uint)