class TestDecisionTree(unittest.TestCase): def setUp(self): self.tree = DecisionTree() self.data = [{"male": True, "tall": False, "rich": True, "married": False}, {"male": False, "tall": True, "rich": True, "married": True}, {"male": False, "tall": True, "rich": False, "married": True}, {"male": True, "tall": True, "rich": True, "married": True}, {"male": False, "tall": True, "rich": False, "married": True}] def test_entropy_corectness(self): entropy = self.tree.entropy(self.data, "married") is_within = entropy > 0.721 and entropy < 0.722 is_within.should.be.ok def test_gain_corectness(self): gain = self.tree.gain(self.data, "male", "married") is_within = gain > 0.321 and gain < 0.322 is_within.should.be.ok def test_find_best_attribute(self): attributes = list(self.data[0].keys()) attributes.remove("married") best_attribute = self.tree.find_best_attribute(self.data, attributes, "married") best_attribute.should.eql("tall") def test_grow_decision_tree(self): attributes = list(self.data[0].keys()) attributes.remove("married") root_node = self.tree.grow(self.data, attributes, "married") root_node.label.should.eql("tall")