import unittest
import networkx as nx
import random

import infomap

def girvan_graphs(zout) :
    """
    Create a graph of 128 vertices, 4 communities, like in
    Community Structure in  social and biological networks.
    Girvan newman, 2002. PNAS June, vol 99 n 12

    Community is node modulo 4
    """

    pout = float(zout)/96.
    pin = (16.-pout*96.)/31.
    graph = nx.Graph()
    graph.add_nodes_from(range(128))
    for x in graph.nodes() :
        for y in graph.nodes() :
            if x < y :
                val = random.random()
                if x % 4 == y % 4 :
                    #nodes belong to the same community
                    if val < pin :
                        graph.add_edge(x, y)

                else :
                    if val < pout :
                        graph.add_edge(x, y)
    return graph


class InfomapTest(unittest.TestCase):

    def test_girvan(self):
        """
        Generate ground truth graph with 4 communities.
        Test that infomap detects them correctly
        """
        graph = girvan_graphs(4)
        modules = infomap.infomap(graph)

        for node, module in modules.items():
            self.assertEqual(module, modules[node%4])

    def test_communites_to_nodes(self):
        """
        Test that the new graph generated in the the second phase has the number of communities as nodes
        """
        g = nx.erdos_renyi_graph(50, 0.1)
        partition = dict([])
        for node in g.nodes() :
            partition[node] = node % 5
        #set up mock object
        info_partition = infomap.Partition(g)
        info_partition.modules = partition

        self.assertSetEqual(set(partition.values()), set(info_partition.second_pass((info_partition.modules)).nodes()))


if __name__ == '__main__':
    unittest.main()
@@ -1387,16 +1387,20 @@
        g
+raph
  = nx.er
@@ -1474,16 +1474,20 @@
 ode in g
+raph
 .nodes()
@@ -1600,16 +1600,20 @@
 tition(g
+raph
 )%0A      
@@ -1771,16 +1771,594 @@
 es()))%0A%0A
+    def test_isolated_nodes(self) :%0A        %22%22%22%0A        Test that the generated graph is isomoprh for the pathological case when all nodes are isolated%0A        %22%22%22%0A        graph = nx.erdos_renyi_graph(50, 0.1)%0A        partition = dict(%5B%5D)%0A        for node in graph.nodes() :%0A            partition%5Bnode%5D = node%0A%0A        #set up mock object%0A        info_partition = infomap.Partition(graph)%0A        info_partition.modules = partition%0A        generated_graph = info_partition.second_pass((info_partition.modules))%0A%0A        self.assertTrue(nx.is_isomorphic(graph, generated_graph))%0A
 %0Aif __na

test for pathological case of isolated nodes