Relationships filter_by_resource_type scope

app/models/relationship.rb

@@ -11,16 +11,27 @@ class Relationship < ApplicationRecord
   # Filtering methods
   #
 
+  def self.filtered(of_type, except_type)
+    relationships = self
+    relationships = relationships.where(:resource_type => of_type) if of_type.present?
+    relationships = relationships.where.not(:resource_type => except_type) if except_type.present?
+    relationships
+  end
+
   def filtered?(of_type, except_type)
     (!of_type.empty? && !of_type.include?(resource_type)) ||
       (!except_type.empty? && except_type.include?(resource_type))
   end
 
   def self.filter_by_resource_type(relationships, options)
-    of_type = options[:of_type].to_miq_a
-    except_type = options[:except_type].to_miq_a
+    of_type = Array.wrap(options[:of_type])
+    except_type = Array.wrap(options[:except_type])
     return relationships if of_type.empty? && except_type.empty?
-    relationships.reject { |r| r.filtered?(of_type, except_type) }
+    if relationships.kind_of?(Array)
+      relationships.reject { |r| r.filtered?(of_type, except_type) }
+    else
+      relationships.filtered(of_type, except_type)
+    end
   end
 
   #

spec/factories/relationship.rb

@@ -1,5 +1,17 @@
 FactoryGirl.define do
-  factory :relationship_vm_vmware, :class => :Relationship do
+  factory :relationship do
     resource_type  "VmOrTemplate"
   end
+
+  factory :relationship_vm_vmware, :parent => :relationship do
+    resource_type  "VmOrTemplate"
+  end
+
+  factory :relationship_host_vmware, :parent => :relationship do
+    resource_type  "Host"
+  end
+
+  factory :relationship_storage_vmware, :parent => :relationship do
+    resource_type  "Storage"
+  end
 end

spec/models/relationship_spec.rb

@@ -1,9 +1,13 @@
 describe Relationship do
-  before(:each) do
-    @rel = FactoryGirl.create(:relationship_vm_vmware)
-  end
+  describe "#filtered?" do
+    before(:each) do
+      @rel = FactoryGirl.build(:relationship_vm_vmware)
+    end
+
+    it "with neither" do
+      expect(@rel).not_to be_filtered([], [])
+    end
 
-  context "#filtered?" do
     it "with of_type" do
       expect(@rel).not_to be_filtered(["VmOrTemplate"], [])
       expect(@rel).not_to be_filtered(["VmOrTemplate", "Host"], [])
@@ -23,4 +27,68 @@
       expect(@rel).to     be_filtered(["Host"], ["VmOrTemplate"])
     end
   end
+
+  describe ".filter_by_resource_type" do
+    let(:storages) { FactoryGirl.build_list(:relationship_storage_vmware, 1) }
+    let(:vms) { FactoryGirl.build_list(:relationship_vm_vmware, 1) }
+    let(:hosts) { FactoryGirl.build_list(:relationship_host_vmware, 1) }
+
+    it "includes" do
+      expect(Relationship.filter_by_resource_type(vms + hosts, :of_type => "Host")).to eq(hosts)
+      expect(Relationship.filter_by_resource_type(vms + hosts, :of_type => ["Host"])).to eq(hosts)
+      expect(Relationship.filter_by_resource_type(vms + hosts, :of_type => ["Host"], :except_type => [])).to eq(hosts)
+    end
+
+    it "includes multi" do
+      expect(Relationship.filter_by_resource_type(vms + hosts + storages, :of_type => %w(Host VmOrTemplate)))
+        .to match_array(vms + hosts)
+    end
+
+    it "includes everything" do
+      expect(Relationship.filter_by_resource_type(vms + hosts, :of_type => %w(Host VmOrTemplate))).to eq(vms + hosts)
+    end
+
+    it "includes nothing" do
+      expect(Relationship.filter_by_resource_type(vms, :of_type => ["Host"])).to be_empty
+    end
+
+    it "excludes" do
+      expect(Relationship.filter_by_resource_type(vms + hosts, :except_type => "Host")).to eq(vms)
+      expect(Relationship.filter_by_resource_type(vms + hosts, :except_type => ["Host"])).to eq(vms)
+      expect(Relationship.filter_by_resource_type(vms + hosts, :except_type => ["Host"], :of_type => [])).to eq(vms)
+    end
+
+    it "excludes multi" do
+      expect(Relationship.filter_by_resource_type(vms + hosts + storages, :except_type => %w(Host VmOrTemplate)))
+        .to eq(storages)
+    end
+
+    it "excludes everything" do
+      expect(Relationship.filter_by_resource_type(vms + hosts, :except_type => %w(Host VmOrTemplate))).to be_empty
+    end
+
+    it "excludes nothing" do
+      expect(Relationship.filter_by_resource_type(vms, :except_type => %w(Host))).to eq(vms)
+    end
+
+    it "includes and excludes" do
+      expect(Relationship.filter_by_resource_type(vms + hosts, :of_type => ["VmOrTemplate"], :except_type => ["Host"]))
+        .to eq(vms)
+    end
+
+    it "neither includes nor excludes" do
+      expect(Relationship.filter_by_resource_type(vms, {})).to eq(vms)
+    end
+
+    it "scopes" do
+      vms.map(&:save!)
+      hosts.map(&:save!)
+      storages.map(&:save!)
+      filtered_results = Relationship.filter_by_resource_type(Relationship.all,
+                                                              :of_type     => %w(Host VmOrTemplate),
+                                                              :except_type => %w(Storage))
+      expect(filtered_results).not_to be_kind_of(Array)
+      expect(filtered_results).to match_array(vms + hosts)
+    end
+  end
 end