From 01428e39728c65c3868c03a53daa82581bdd86b4 Mon Sep 17 00:00:00 2001
From: Eryn Wells <eryn@erynwells.me>
Date: Sun, 8 Oct 2017 07:23:03 -0700
Subject: [PATCH] Fix squares; add dimension property; add some tests

---
 sudoku.py | 21 ++++++++++++++++-----
 test.py   | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 5 deletions(-)
 create mode 100644 test.py

diff --git a/sudoku.py b/sudoku.py
index fc8411a..8f05749 100644
--- a/sudoku.py
+++ b/sudoku.py
@@ -12,10 +12,14 @@ class Sudoku:
         dim = math.sqrt(size)
         if dim != int(dim):
             raise ValueError('Board size must have an integral square root.')
-        self.dimension = int(dim)
+        self._dimension = int(dim)
         self.size = size
         self.board = bytearray(b'\x00' * (size * size))
 
+    @property
+    def dimension(self):
+        return self._dimension
+
     @property
     def rows(self):
         sz = self.size
@@ -27,6 +31,11 @@ class Sudoku:
         sz2 = sz ** 2
         return [range(i, sz2, sz) for i in range(sz)]
 
+    @property
+    def squares(self):
+        dim = self.dimension
+        return [self.square(x, y) for y in range(dim) for x in range(dim)]
+
     @property
     def solved(self):
         def _check_group(group):
@@ -40,14 +49,16 @@ class Sudoku:
         # TODO: WIP
 
     def square(self, x, y):
-        dim = int(math.sqrt(self.size))
+        dim = self.dimension
         if (x < 0 or x >= dim) or (y < 0 or y >= dim):
             raise IndexError('Invalid coordinates for square: ({}, {})'.format(x, y))
-        offset_x = x * dim
-        offset_y = y * dim
+
+        offset = (x * dim, y * dim * self.size)
+
         def _range(i):
-            start = offset_x + i * self.size
+            start = (offset[1] + i * self.size) + offset[0]
             return range(start, start + dim)
+
         ranges = itertools.chain(*[_range(i) for i in range(dim)])
         return ranges
 
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..38805d6
--- /dev/null
+++ b/test.py
@@ -0,0 +1,53 @@
+#!env python3
+# Eryn Wells <eryn@erynwells.me>
+'''
+Unit tests for the Sudoku module.
+'''
+
+import unittest
+import sudoku
+
+class Sudoku4Tests(unittest.TestCase):
+    def setUp(self):
+        self.board = sudoku.Sudoku(size=4)
+
+    def test_that_board_is_sane(self):
+        self.assertEqual(self.board.size, 4)
+        self.assertEqual(len(self.board.board), 4**2)
+        self.assertEqual(self.board.dimension, 2)
+
+    def test_rows(self):
+        expected_rows = [
+            [ 0,  1,  2,  3],
+            [ 4,  5,  6,  7],
+            [ 8,  9, 10, 11],
+            [12, 13, 14, 15]
+        ]
+        for (row, exrow) in zip(self.board.rows, expected_rows):
+            row_list = list(row)
+            with self.subTest(row=row_list, ex=exrow):
+                self.assertEqual(row_list, exrow)
+
+    def test_columns(self):
+        expected_columns = [
+            [0, 4,  8, 12],
+            [1, 5,  9, 13],
+            [2, 6, 10, 14],
+            [3, 7, 11, 15]
+        ]
+        for (col, excol) in zip(self.board.columns, expected_columns):
+            col_list = list(col)
+            with self.subTest(col=col_list, ex=excol):
+                self.assertEqual(col_list, excol)
+
+    def test_squares(self):
+        expected_squares = [
+            [ 0,  1,  4,  5],
+            [ 2,  3,  6,  7],
+            [ 8,  9, 12, 13],
+            [10, 11, 14, 15]
+        ]
+        for (sq, exsq) in zip(self.board.squares, expected_squares):
+            sq_list = list(sq)
+            with self.subTest(sq=sq_list, ex=exsq):
+                self.assertEqual(sq_list, exsq)