test_ml.py

from algorithms.ml.nearest_neighbor import (
    distance,
    nearest_neighbor
)

import unittest


class TestML(unittest.TestCase):
    def setUp(self):
        # train set for the AND-function
        self.trainSetAND = {(0, 0): 0, (0, 1): 0, (1, 0): 0, (1, 1): 1}

        # train set for light or dark colors
        self.trainSetLight = {(11, 98, 237): 'L', (3, 39, 96): 'D',
                              (242, 226, 12): 'L', (99, 93, 4): 'D',
                              (232, 62, 32): 'L', (119, 28, 11): 'D',
                              (25, 214, 47): 'L', (89, 136, 247): 'L',
                              (21, 34, 63): 'D', (237, 99, 120): 'L',
                              (73, 33, 39): 'D'}

    def test_nearest_neighbor(self):
        # AND-function
        self.assertEqual(nearest_neighbor((1, 1), self.trainSetAND), 1)
        self.assertEqual(nearest_neighbor((0, 1), self.trainSetAND), 0)

        # dark/light color test
        self.assertEqual(nearest_neighbor((31, 242, 164),
                                          self.trainSetLight), 'L')
        self.assertEqual(nearest_neighbor((13, 94, 64),
                                          self.trainSetLight), 'D')
        self.assertEqual(nearest_neighbor((230, 52, 239),
                                          self.trainSetLight), 'L')

    def test_distance(self):
        self.assertAlmostEqual(distance((1, 2, 3), (1, 0, -1)), 4.47, 2)


if __name__ == "__main__":
    unittest.main()