From 02ac35df9206ceaf9469ce5e77555b4809c1abcf Mon Sep 17 00:00:00 2001
From: Jeremy Thurgood <firxen@gmail.com>
Date: Thu, 15 May 2014 20:23:56 +0200
Subject: [PATCH] Better location shifting.

---
 naja/gameboard.py            | 55 +++++++++++++-----------------------
 naja/tests/test_gameboard.py | 39 +++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 36 deletions(-)

diff --git a/naja/gameboard.py b/naja/gameboard.py
index d626849..edda43c 100644
--- a/naja/gameboard.py
+++ b/naja/gameboard.py
@@ -104,48 +104,31 @@ class GameBoard(object):
         location = LocationCard.new_location(choice(self.locations).copy())
         self.board_locations[position] = location
 
-    def shift_locations(self, direction):
+    def shift_location_row(self, change, is_vertical):
         px, py = self.player.position
         shifted_locations = {}
-        # TODO: Make this less horrible. Also test it.
+        mkpos = lambda i: (px, i) if is_vertical else (i, py)
+
+        for i in range(5):
+            if (px, py) == mkpos(i):
+                continue
+            new_i = (i + change) % 5
+            if (px, py) == mkpos(new_i):
+                new_i = (new_i + change) % 5
+            shifted_locations[mkpos(new_i)] = self.board_locations[mkpos(i)]
+
+        print change, is_vertical, shifted_locations
+        self.board_locations.update(shifted_locations)
+
+    def shift_locations(self, direction):
         if BITS[direction] == BITS.NORTH:
-            for y in range(5):
-                if y == py:
-                    continue
-                new_y = y - 1
-                if new_y == py:
-                    new_y -= 1
-                new_y %= 5
-                shifted_locations[(px, new_y)] = self.board_locations[(px, y)]
+            self.shift_location_row(-1, is_vertical=True)
         elif BITS[direction] == BITS.SOUTH:
-            for y in range(5):
-                if y == py:
-                    continue
-                new_y = y + 1
-                if new_y == py:
-                    new_y += 1
-                new_y %= 5
-                shifted_locations[(px, new_y)] = self.board_locations[(px, y)]
+            self.shift_location_row(1, is_vertical=True)
         elif BITS[direction] == BITS.EAST:
-            for x in range(5):
-                if x == px:
-                    continue
-                new_x = x + 1
-                if new_x == px:
-                    new_x += 1
-                new_x %= 5
-                shifted_locations[(new_x, py)] = self.board_locations[(x, py)]
+            self.shift_location_row(1, is_vertical=False)
         elif BITS[direction] == BITS.WEST:
-            for x in range(5):
-                if x == px:
-                    continue
-                new_x = x - 1
-                if new_x == px:
-                    new_x -= 1
-                new_x %= 5
-                shifted_locations[(new_x, py)] = self.board_locations[(x, py)]
-
-        self.board_locations.update(shifted_locations)
+            self.shift_location_row(-1, is_vertical=False)
 
     def change_mode(self, new_mode):
         """Advance to the next mode"""
diff --git a/naja/tests/test_gameboard.py b/naja/tests/test_gameboard.py
index bbc5128..5dae954 100644
--- a/naja/tests/test_gameboard.py
+++ b/naja/tests/test_gameboard.py
@@ -85,6 +85,45 @@ class TestGameBoard(TestCase):
 
         self.assert_state(state_1, state_2)
 
+    def generate_locations(self, override_dict=None):
+        locations_dict = dict(((x, y), '%s%s' % (x, y))
+                              for x in range(5) for y in range(5))
+        if override_dict:
+            locations_dict.update(override_dict)
+        return locations_dict
+
+    def test_shift_locations_north(self):
+        board = GameBoard.new_game([{'actions': []}])
+        board.board_locations = self.generate_locations()
+        board.shift_locations('NORTH')
+        self.assertEqual(board.board_locations, self.generate_locations({
+            (2, 0): '21', (2, 1): '23', (2, 3): '24', (2, 4): '20',
+        }))
+
+    def test_shift_locations_south(self):
+        board = GameBoard.new_game([{'actions': []}])
+        board.board_locations = self.generate_locations()
+        board.shift_locations('SOUTH')
+        self.assertEqual(board.board_locations, self.generate_locations({
+            (2, 0): '24', (2, 1): '20', (2, 3): '21', (2, 4): '23',
+        }))
+
+    def test_shift_locations_east(self):
+        board = GameBoard.new_game([{'actions': []}])
+        board.board_locations = self.generate_locations()
+        board.shift_locations('EAST')
+        self.assertEqual(board.board_locations, self.generate_locations({
+            (0, 2): '42', (1, 2): '02', (3, 2): '12', (4, 2): '32',
+        }))
+
+    def test_shift_locations_west(self):
+        board = GameBoard.new_game([{'actions': []}])
+        board.board_locations = self.generate_locations()
+        board.shift_locations('WEST')
+        self.assertEqual(board.board_locations, self.generate_locations({
+            (0, 2): '12', (1, 2): '32', (3, 2): '42', (4, 2): '02',
+        }))
+
 
 class TestLocationCard(TestCase):
     def test_generate_bitwise_operand(self):
-- 
2.34.1