<commit_before>import unittest
import BayesianModel as bm
import networkx as nx


class TestModel(unittest.TestCase):

    def setUp(self):
        self.G = bm.BayesianModel()

    def test_class_init(self):
        self.assertIsInstance(self.G, nx.DiGraph)

    def test_add_nodes_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.assertSameElements(self.G.nodes(), ['a', 'b', 'c', 'd'])

    def test_add_nodes_non_string(self):
        self.assertRaises(TypeError, self.G.add_nodes, [1, 2, 3, 4])

    def test_add_edges_both_tuples(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges(('a', 'b'), ('c', 'd'))
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('a', 'd'),
                                                 ('b', 'c'), ('b', 'd')])

    def test_add_edges_tail_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges('a', ('b', 'c'))
        self.assertSameElements(self.G.edges(), [('a', 'b'), ('a', 'c')])

    def test_add_edges_head_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges(('a', 'b'), 'c')
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('b', 'c')])

    def test_add_edges_both_string(self):
        self.G.add_nodes('a', 'b')
        self.G.add_edges('a', 'b')
        self.assertSameElements(self.G.edges(), [('a', 'b')])

    def test_add_edges_multiple_times(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges('a', ('c', 'd'))
        self.G.add_edges('b', ('c', 'd'))
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('a', 'd'),
                                                 ('b', 'c'), ('b', 'd')])

    def tearDown(self):
        del self.G


class TestNodeProperties(unittest.TestCase):

    def setUp(self):
        self.G = bm.BayesianModel()
        self.G.add_nodes('a', 'b', 'c', 'd', 'e')
        self.G.add_edges(('a', 'b'), ('c', 'd'))

    def test_parents(self):
        self.assertSameElements(self.G.node['c']['_parents'], ['a', 'b'])
        self.assertSameElements(self.G.node['d']['_parents'], ['a', 'b'])
# TODO       self.assertRaises(KeyError, self.G.node['a']['_parents'])
# TODO       self.assertRaises(KeyError, self.G.node['b']['_parents'])

# TODO add test_default_rule

    def test_add_states(self):
        self.G.add_states('a', ('test_state_3', 'test_state_1',
                                'test_state_2'))
        self.G.add_states('b', ('test_state_2', 'test_state_1'))
        self.G.add_states('c', ('test_state_1',))
        self.G.add_states('d', ('test_state_1', 'test_state_2'))
        self.assertEqual(self.G.node['a']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0],
                                                       ['test_state_3', 0]])
        self.assertEqual(self.G.node['b']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0]])
        self.assertEqual(self.G.node['c']['_states'], [['test_state_1', 0]])
        self.assertEqual(self.G.node['d']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0]])

# TODO add test_default_rule_for_states

# TODO    def test_rule_for_states

    def test_states_function(self):
        self.G.add_states('a', ('test_state_3', 'test_state_1',
                                'test_state_2'))
        self.G.add_states('b', ('test_state_2', 'test_state_1'))
        self.G.add_states('c', ('test_state_1',))
        self.G.add_states('d', ('test_state_1', 'test_state_2'))
        states = {'a': [], 'b': [], 'c': [], 'd': []}
        nodes = ['a', 'b', 'c', 'd']
        for node in nodes:
            for state in self.G.states(node):
                states[node].append(state)
        self.assertEqual(states['a'], ['test_state_1', 'test_state_2',
                                       'test_state_3'])
        self.assertEqual(states['b'], ['test_state_1', 'test_state_2'])
        self.assertEqual(states['c'], ['test_state_1'])
        self.assertEqual(states['d'], ['test_state_1', 'test_state_2'])

    def tearDown(self):
        del self.G

if __name__ == '__main__':
        unittest.main()<commit_msg>replaced add_states function with set_states<commit_after>import unittest
import BayesianModel as bm
import networkx as nx


class TestModel(unittest.TestCase):

    def setUp(self):
        self.G = bm.BayesianModel()

    def test_class_init(self):
        self.assertIsInstance(self.G, nx.DiGraph)

    def test_add_nodes_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.assertSameElements(self.G.nodes(), ['a', 'b', 'c', 'd'])

    def test_add_nodes_non_string(self):
        self.assertRaises(TypeError, self.G.add_nodes, [1, 2, 3, 4])

    def test_add_edges_both_tuples(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges(('a', 'b'), ('c', 'd'))
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('a', 'd'),
                                                 ('b', 'c'), ('b', 'd')])

    def test_add_edges_tail_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges('a', ('b', 'c'))
        self.assertSameElements(self.G.edges(), [('a', 'b'), ('a', 'c')])

    def test_add_edges_head_string(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges(('a', 'b'), 'c')
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('b', 'c')])

    def test_add_edges_both_string(self):
        self.G.add_nodes('a', 'b')
        self.G.add_edges('a', 'b')
        self.assertSameElements(self.G.edges(), [('a', 'b')])

    def test_add_edges_multiple_times(self):
        self.G.add_nodes('a', 'b', 'c', 'd')
        self.G.add_edges('a', ('c', 'd'))
        self.G.add_edges('b', ('c', 'd'))
        self.assertSameElements(self.G.edges(), [('a', 'c'), ('a', 'd'),
                                                 ('b', 'c'), ('b', 'd')])

    def tearDown(self):
        del self.G


class TestNodeProperties(unittest.TestCase):

    def setUp(self):
        self.G = bm.BayesianModel()
        self.G.add_nodes('a', 'b', 'c', 'd', 'e')
        self.G.add_edges(('a', 'b'), ('c', 'd'))

    def test_parents(self):
        self.assertSameElements(self.G.node['c']['_parents'], ['a', 'b'])
        self.assertSameElements(self.G.node['d']['_parents'], ['a', 'b'])
# TODO       self.assertRaises(KeyError, self.G.node['a']['_parents'])
# TODO       self.assertRaises(KeyError, self.G.node['b']['_parents'])

# TODO add test_default_rule

    def test_set_states(self):
        self.G.set_states('a', ('test_state_3', 'test_state_1',
                                'test_state_2'))
        self.G.set_states('b', ('test_state_2', 'test_state_1'))
        self.G.set_states('c', ('test_state_1',))
        self.G.set_states('d', ('test_state_1', 'test_state_2'))
        self.assertEqual(self.G.node['a']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0],
                                                       ['test_state_3', 0]])
        self.assertEqual(self.G.node['b']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0]])
        self.assertEqual(self.G.node['c']['_states'], [['test_state_1', 0]])
        self.assertEqual(self.G.node['d']['_states'], [['test_state_1', 0],
                                                       ['test_state_2', 0]])

# TODO add test_default_rule_for_states

# TODO    def test_rule_for_states

    def test_states_function(self):
        self.G.set_states('a', ('test_state_3', 'test_state_1',
                                'test_state_2'))
        self.G.set_states('b', ('test_state_2', 'test_state_1'))
        self.G.set_states('c', ('test_state_1',))
        self.G.set_states('d', ('test_state_1', 'test_state_2'))
        states = {'a': [], 'b': [], 'c': [], 'd': []}
        nodes = ['a', 'b', 'c', 'd']
        for node in nodes:
            for state in self.G.states(node):
                states[node].append(state)
        self.assertEqual(states['a'], ['test_state_1', 'test_state_2',
                                       'test_state_3'])
        self.assertEqual(states['b'], ['test_state_1', 'test_state_2'])
        self.assertEqual(states['c'], ['test_state_1'])
        self.assertEqual(states['d'], ['test_state_1', 'test_state_2'])

    def tearDown(self):
        del self.G

if __name__ == '__main__':
        unittest.main()
{'fraction_non_alphanumeric': 0.20355111987930835, 'fraction_numerical': 0.01021237089474295, 'mean_word_length': 2.3158907272027704, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5496727', 'n_tokens_mistral': 2989, 'n_tokens_neox': 2929, 'n_words': 493}