#!/usr/bin/env python
from nose.tools import assert_equal
import os
import tempfile

import networkx as nx
from networkx.testing.utils import *


class TestGpickle(object):
    def setUp(self):
        G=nx.Graph(name="test")
        e=[('a','b'),('b','c'),('c','d'),('d','e'),('e','f'),('a','f')]
        G.add_edges_from(e,width=10)
        G.add_node('g',color='green')
        G.graph['number']=1
        DG=nx.DiGraph(G)
        MG=nx.MultiGraph(G)
        MG.add_edge('a', 'a')
        MDG=nx.MultiDiGraph(G)
        MDG.add_edge('a', 'a')
        fG = G.copy()
        fDG = DG.copy()
        fMG = MG.copy()
        fMDG = MDG.copy()
        nx.freeze(fG)
        nx.freeze(fDG)
        nx.freeze(fMG)
        nx.freeze(fMDG)
        self.G=G
        self.DG=DG
        self.MG=MG
        self.MDG=MDG
        self.fG=fG
        self.fDG=fDG
        self.fMG=fMG
        self.fMDG=fMDG

    def test_gpickle(self):
        for G in [self.G, self.DG, self.MG, self.MDG,
                  self.fG, self.fDG, self.fMG, self.fMDG]:
            (fd,fname)=tempfile.mkstemp()
            nx.write_gpickle(G,fname)
            Gin=nx.read_gpickle(fname)
            assert_nodes_equal(list(G.nodes(data=True)),
                               list(Gin.nodes(data=True)))
            assert_edges_equal(list(G.edges(data=True)),
                               list(Gin.edges(data=True)))
            assert_graphs_equal(G, Gin)
            os.close(fd)
            os.unlink(fname)

    def test_protocol(self):
        for G in [self.G, self.DG, self.MG, self.MDG,
                  self.fG, self.fDG, self.fMG, self.fMDG]:
            with tempfile.TemporaryFile() as f:
                nx.write_gpickle(G, f, 0)
                f.seek(0)
                Gin = nx.read_gpickle(f)
                assert_nodes_equal(list(G.nodes(data=True)),
                                   list(Gin.nodes(data=True)))
                assert_edges_equal(list(G.edges(data=True)),
                                   list(Gin.edges(data=True)))
                assert_graphs_equal(G, Gin)