lib/thor/spree/bootstrap_cleanup.thor in datashift-0.6.0 vs lib/thor/spree/bootstrap_cleanup.thor in datashift-0.6.1

@@ -14,29 +14,39 @@
   
     include DataShift::Logging
        
     desc "cleanup", "Remove Spree Product/Variant data from DB"
     
-    def cleanup()  #, [:input, :verbose, :sku_prefix] => :environment do |t, args|
+    def cleanup()
 
-     require 'spree_helper'
+      require 'spree_helper'
 
       require File.expand_path('config/environment.rb')
-    
-      %w{Image OptionType OptionValue Product Property ProductGroup ProductProperty ProductOptionType Variant Taxonomy Taxon Zone}.each do |k|
-        instance_variable_set("@#{k}_klass", DataShift::SpreeHelper::get_spree_class(k)) 
-        puts "Clearing model #{DataShift::SpreeHelper::get_spree_class(k)}"
-        instance_variable_get("@#{k}_klass").delete_all
+
+      cleanup =  %w{ Image OptionType OptionValue 
+                    Product Property ProductGroup ProductProperty ProductOptionType 
+                    Variant Taxonomy Taxon Zone
+      }
+
+      cleanup.each do |k|
+        klass = DataShift::SpreeHelper::get_spree_class(k)
+      if(klass)
+        puts "Clearing model #{klass}"
+        klass.delete_all
+      else
+        puts "WARNING - Coulod not find AR model for class name #{k}"
       end
+        
+    end
       
-      if(File.exists?('public/spree/products') )
-        puts "Removing old Product assets from 'public/spree/products'"
-        FileUtils::rm_rf('public/public/spree/products') 
-      end
+    if(File.exists?('public/spree/products') )
+      puts "Removing old Product assets from 'public/spree/products'"
+      FileUtils::rm_rf('public/public/spree/products') 
+    end
       
-      FileUtils::rm_rf('MissingRecords') if(File.exists?('MissingRecords') )
+    FileUtils::rm_rf('MissingRecords') if(File.exists?('MissingRecords') )
       
-    end
-  
   end
+  
+end
 
 end
\ No newline at end of file