scenarios can contain control links

app/models/scenario_import.rb

@@ -60,6 +60,7 @@ class ScenarioImport
     description = parsed_data['description']
     name = parsed_data['name']
     links = parsed_data['links']
+    control_links = parsed_data['control_links'] || []
     tag_fg_color = parsed_data['tag_fg_color']
     tag_bg_color = parsed_data['tag_bg_color']
     source_url = parsed_data['source_url'].presence || nil
@@ -87,12 +88,19 @@ class ScenarioImport
         end
         agent
       end
+
       if success
         links.each do |link|
           receiver = created_agents[link['receiver']]
           source = created_agents[link['source']]
           receiver.sources << source unless receiver.sources.include?(source)
         end
+
+        control_links.each do |control_link|
+          controller = created_agents[control_link['controller']]
+          control_target = created_agents[control_link['control_target']]
+          controller.control_targets << control_target unless controller.control_targets.include?(control_target)
+        end
       end
     end
 

lib/agents_exporter.rb

@@ -20,7 +20,8 @@ class AgentsExporter
       :tag_bg_color => options[:tag_bg_color],
       :exported_at => Time.now.utc.iso8601,
       :agents => agents.map { |agent| agent_as_json(agent) },
-      :links => links
+      :links => links,
+      :control_links => control_links
     }
   end
 
@@ -32,14 +33,26 @@ class AgentsExporter
     agent_ids = agents.map(&:id)
 
     contained_links = agents.map.with_index do |agent, index|
-      agent.links_as_source.where(:receiver_id => agent_ids).map do |link|
-        { :source => index, :receiver => agent_ids.index(link.receiver_id) }
+      agent.links_as_source.where(receiver_id: agent_ids).map do |link|
+        { source: index, receiver: agent_ids.index(link.receiver_id) }
       end
     end
 
     contained_links.flatten.compact
   end
 
+  def control_links
+    agent_ids = agents.map(&:id)
+
+    contained_controller_links = agents.map.with_index do |agent, index|
+      agent.control_links_as_controller.where(control_target_id: agent_ids).map do |control_link|
+        { controller: index, control_target: agent_ids.index(control_link.control_target_id) }
+      end
+    end
+
+    contained_controller_links.flatten.compact
+  end
+
   def agent_as_json(agent)
     {
       :type => agent.type,

spec/lib/agents_exporter_spec.rb

@@ -25,6 +25,7 @@ describe AgentsExporter do
       expect(data[:tag_bg_color]).to eq(tag_bg_color)
       expect(Time.parse(data[:exported_at])).to be_within(2).of(Time.now.utc)
       expect(data[:links]).to eq([{ :source => 0, :receiver => 1 }])
+      expect(data[:control_links]).to eq([])
       expect(data[:agents]).to eq(agent_list.map { |agent| exporter.agent_as_json(agent) })
       expect(data[:agents].all? { |agent_json| agent_json[:guid].present? && agent_json[:type].present? && agent_json[:name].present? }).to be_truthy
 
@@ -38,6 +39,13 @@ describe AgentsExporter do
 
       expect(exporter.as_json[:links]).to eq([{ :source => 0, :receiver => 1 }])
     end
+
+    it "outputs control links to agents within the incoming set, but not outside it" do
+      agents(:jane_rain_notifier_agent).control_targets = [agents(:jane_weather_agent), agents(:jane_basecamp_agent)]
+      agents(:jane_rain_notifier_agent).save!
+
+      expect(exporter.as_json[:control_links]).to eq([{ :controller => 1, :control_target => 0 }])
+    end
   end
 
   describe "#filename" do

spec/models/scenario_import_spec.rb

@@ -74,7 +74,8 @@ describe ScenarioImport do
       ],
       :links => [
         { :source => 0, :receiver => 1 }
-      ]
+      ],
+      :control_links => []
     }
   end
   let(:valid_data) { valid_parsed_data.to_json }
@@ -226,6 +227,38 @@ describe ScenarioImport do
             scenario_import.import
           }.to change { users(:bob).agents.count }.by(2)
         end
+
+        describe "with control links" do
+          it 'creates the links' do
+            valid_parsed_data[:control_links] = [
+              { :controller => 1, :control_target => 0 }
+            ]
+
+            expect {
+              scenario_import.import
+            }.to change { users(:bob).agents.count }.by(2)
+
+            weather_agent = scenario_import.scenario.agents.find_by(:guid => "a-weather-agent")
+            trigger_agent = scenario_import.scenario.agents.find_by(:guid => "a-trigger-agent")
+
+            expect(trigger_agent.sources).to eq([weather_agent])
+            expect(weather_agent.controllers.to_a).to eq([trigger_agent])
+            expect(trigger_agent.control_targets.to_a).to eq([weather_agent])
+          end
+
+          it "doesn't crash without any control links" do
+            valid_parsed_data.delete(:control_links)
+
+            expect {
+              scenario_import.import
+            }.to change { users(:bob).agents.count }.by(2)
+
+            weather_agent = scenario_import.scenario.agents.find_by(:guid => "a-weather-agent")
+            trigger_agent = scenario_import.scenario.agents.find_by(:guid => "a-trigger-agent")
+
+            expect(trigger_agent.sources).to eq([weather_agent])
+          end
+        end
       end
 
       describe "#generate_diff" do