lib/graph.rb in aws_security_viz-0.1.7.pre.alpha.pre.145 vs lib/graph.rb in aws_security_viz-0.2.0.pre.alpha.pre.152

@@ -5,15 +5,17 @@
 
   def initialize(config, underlying=RGL::DirectedAdjacencyGraph.new)
     @config = config
     @underlying = underlying
     @edge_properties = {}
+    @node_properties = {}
   end
 
-  def add_node(name)
-    log("node: #{name}")
+  def add_node(name, opts)
+    log("node: #{name}, opts: #{opts}")
     @underlying.add_vertex(name)
+    @node_properties[name] = opts
   end
 
   def add_edge(from, to, opts)
     log("edge: #{from} -> #{to}")
     @underlying.add_edge(from, to)
@@ -23,10 +25,10 @@
   def filter(source, destination)
     @underlying = GraphFilter.new(underlying).filter(source, destination)
   end
 
   def output(renderer)
-    @underlying.each_vertex { |v| renderer.add_node(v) }
+    @underlying.each_vertex { |v| renderer.add_node(v, @node_properties[v] || {}) }
     @underlying.each_edge { |u, v|
       renderer.add_edge(u, v, opts(u, v))
     }
     renderer.output
   end