import json

from thug.Logging.modules.ExploitGraph import ExploitGraph


class TestExploitGraph:
    url    = "www.example.com"
    source = "www.ex1.com"
    dest   = "www.ex2.com"
    method = "iframe"

    exploit_graph = ExploitGraph(url)

    def test_init(self):
        assert self.exploit_graph.G.graph["url"] in (self.url, )
        assert self.exploit_graph.G.graph["name"] in ("thug-exploit-graph", )

    def test_add_connection(self):
        self.exploit_graph.add_connection(self.source, self.dest, self.method)
        nodes = self.exploit_graph.G.nodes
        assert self.source in nodes
        assert self.dest in nodes

        edge = list(self.exploit_graph.G.edges(data = True))[0]
        assert self.source in edge
        assert self.dest in edge
        assert {'method': self.method} in edge

    def test_draw(self):
        graph = json.loads(self.exploit_graph.draw())
        assert graph["directed"]
        assert graph["multigraph"]
        assert graph["graph"]
        assert graph["nodes"]
        assert graph["links"]
