diff --git a/Gemfile b/Gemfile
index 90b87fc088f4c6a16f9f2c8db9bd987fd99f0ed1..6ad25823f8ae7595d778aad063bcbbc912590094 100644
--- a/Gemfile
+++ b/Gemfile
@@ -1,6 +1,11 @@
-source "https://rubygems.org"
+# frozen_string_literal: true
 
-git_source(:github) {|repo_name| "https://github.com/#{repo_name}" }
+source 'https://rubygems.org'
 
 # Specify your gem's dependencies in foreman_wds.gemspec
 gemspec
+
+gem 'minitest', '~> 5'
+gem 'rake', '~> 13'
+gem 'rubocop', '~> 1'
+gem 'simplecov', '~> 0'
diff --git a/Rakefile b/Rakefile
index d433a1edc1b09f60e4966db4839d9c57b2828d98..dc4712ac605d89215609793d741cc04f4066ab41 100644
--- a/Rakefile
+++ b/Rakefile
@@ -1,10 +1,16 @@
-require "bundler/gem_tasks"
-require "rake/testtask"
+# frozen_string_literal: true
+
+require 'bundler/gem_tasks'
+require 'rake/testtask'
 
 Rake::TestTask.new(:test) do |t|
-  t.libs << "test"
-  t.libs << "lib"
-  t.test_files = FileList["test/**/*_test.rb"]
+  t.libs << 'test'
+  t.libs << 'lib'
+  t.test_files = FileList['test/**/*_test.rb']
 end
 
-task :default => :test
+require 'rubocop/rake_task'
+
+RuboCop::RakeTask.new
+
+task default: :test
diff --git a/app/assets/javascripts/foreman_wds/host_edit_extensions.js b/app/assets/javascripts/foreman_wds/host_edit_extensions.js
index f6e5a67938a6ea387f7d5a1af040768d7312d64d..19c224e9fed1ffb472f953caa9a09aeee3847f5e 100644
--- a/app/assets/javascripts/foreman_wds/host_edit_extensions.js
+++ b/app/assets/javascripts/foreman_wds/host_edit_extensions.js
@@ -1,4 +1,4 @@
-function wds_server_selected(element){
+function wds_server_selected(element) {
   var url = $(element).attr('data-url');
   var type = $(element).attr('data-type');
   var attrs = {};
@@ -17,10 +17,7 @@ function wds_server_selected(element){
   });
 }
 
-var old_os_selected = os_selected;
-os_selected = function(element){
-  old_os_selected(element);
-
+function wds_os_selected() {
   if ($('#os_select select').val() === '') {
     $('#wds_server_select select').val('');
     $('#wds_image_select select').val('');
@@ -35,11 +32,10 @@ os_selected = function(element){
   }
 };
 
-var old_onHostEditLoad = onHostEditLoad;
-onHostEditLoad = function() {
-  old_onHostEditLoad();
-
+function wds_content_loaded() {
   $('#wds_provisioning').detach().insertBefore('#media_select');
+  $('#host_provision_method_build').prop('disabled', false);
+  $('#host_provision_method_wds').prop('disabled', false);
 };
 
 
@@ -52,7 +48,10 @@ function wds_provision_method_selected() {
     $('#wds_image_select select').attr('disabled', true);
   }
 }
-$(document).on('change', '#host_provision_method_wds', wds_provision_method_selected);
+$(document)
+  .on('change', '#host_provision_method_wds', wds_provision_method_selected)
+  .on('change', '.host-architecture-os-select', wds_os_selected)
+  .on('ContentLoad', wds_content_loaded);
 
 $(function() {
   if($('#host_provision_method_wds').is(':checked')) {
diff --git a/app/controllers/concerns/foreman/controller/parameters/wds_server.rb b/app/controllers/concerns/foreman/controller/parameters/wds_server.rb
index 62a39242db8763b85b8339739777b34f420dd00f..31b9475f0fb4e92e4c86bacb5d1381bca214fc57 100644
--- a/app/controllers/concerns/foreman/controller/parameters/wds_server.rb
+++ b/app/controllers/concerns/foreman/controller/parameters/wds_server.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Foreman::Controller::Parameters::WdsServer
   extend ActiveSupport::Concern
 
diff --git a/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb b/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
index dd9897a327ed7abc48f99ce89f3a6760eca2e0dd..c9e53511a553fbb4d380e095b228e0fb9acb4da2 100644
--- a/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
+++ b/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
@@ -1,5 +1,13 @@
+# frozen_string_literal: true
+
 module ForemanWds
   module HostsControllerExtensions
+    included do
+      before_action :cleanup_wds_params
+
+      define_action_permission %w[wds_server_selected], :edit
+    end
+
     def wds_server_selected
       host = @host || item_object
       wds_facet = host.wds_facet || host.build_wds_facet
@@ -8,20 +16,9 @@ module ForemanWds
       render partial: 'wds_servers/image_select', locals: { item: wds_facet }
     end
 
-    def host_params(top_level_hash = controller_name.singularize)
+    def cleanup_wds_params
       # Don't create a WDS facet unless provisioning with it
       params[:host].delete :wds_facet_attributes if params[:host] && params[:host][:provision_method] != 'wds'
-
-      super(top_level_hash)
-    end
-
-    def action_permission
-      case params[:action]
-      when 'wds_server_selected', 'wds_image_selected'
-        :edit
-      else
-        super
-      end
     end
   end
 end
diff --git a/app/controllers/concerns/foreman_wds/unattended_controller_extensions.rb b/app/controllers/concerns/foreman_wds/unattended_controller_extensions.rb
index 7f4633b9daff10a6e19cdea2256ae6b28be60609..143f7062ab581eae9f7a50cabccd913914ba6767 100644
--- a/app/controllers/concerns/foreman_wds/unattended_controller_extensions.rb
+++ b/app/controllers/concerns/foreman_wds/unattended_controller_extensions.rb
@@ -1,8 +1,11 @@
+# frozen_string_literal: true
+
 module ForemanWds
   module UnattendedControllerExtensions
     def host_template
       return wds_render_csr if params[:kind] == 'csr_attributes'
       return wds_deploy_localboot if params[:kind] == 'wds_localboot'
+
       super
     end
 
@@ -35,8 +38,9 @@ module ForemanWds
 
       render inline: "Success. Local boot template was deployed successfully.\n"
     rescue StandardError => e
-      message = format('Failed to set local boot template: %{error}', error: e)
+      message = format('Failed to set local boot template: %<error>s', error: e)
       logger.error message
+
       render text: message, status: :error, content_type: 'text/plain'
     end
   end
diff --git a/app/controllers/wds_servers_controller.rb b/app/controllers/wds_servers_controller.rb
index 7d38ee69367d9037ada7fa64ea3a1690b8b9d86e..81f6540811163100f69ab28b787ed456169054b2 100644
--- a/app/controllers/wds_servers_controller.rb
+++ b/app/controllers/wds_servers_controller.rb
@@ -1,9 +1,15 @@
-class WdsServersController < ::ApplicationController
+# frozen_string_literal: true
+
+class WdsServersController < ApplicationController
   include Foreman::Controller::AutoCompleteSearch
   include Foreman::Controller::Parameters::WdsServer
 
   before_action :find_server, except: %i[index new create]
 
+  def model_of_controller
+    WdsServer
+  end
+
   def index
     @wds_servers = resource_base_search_and_page
   end
diff --git a/app/lib/foreman_wds/wds_boot_image.rb b/app/lib/foreman_wds/wds_boot_image.rb
index bbe9dfd06b2678340556e0e2e9f50727ceb73e67..bac9a37c354dba3a97327df33b406164b79935b2 100644
--- a/app/lib/foreman_wds/wds_boot_image.rb
+++ b/app/lib/foreman_wds/wds_boot_image.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class ForemanWds::WdsBootImage < ForemanWds::WdsImage
   def initialize(json = {})
     super json
@@ -5,6 +7,7 @@ class ForemanWds::WdsBootImage < ForemanWds::WdsImage
 
   def reload
     return false if wds_server.nil?
+
     @json = wds_server.boot_image(name)
     load!
   end
diff --git a/app/lib/foreman_wds/wds_image.rb b/app/lib/foreman_wds/wds_image.rb
index 2e8f6d5ac7a9bb721133602beaa50b393924b4ff..a473fcf74de537fcb795a117d6804d7f5e2d73be 100644
--- a/app/lib/foreman_wds/wds_image.rb
+++ b/app/lib/foreman_wds/wds_image.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class ForemanWds::WdsImage
   WDS_IMAGE_ARCHES = [nil, /^i.86|x86$/i, /^ia64$/i, /^x86_64|x64$/i, /^arm$/i].freeze
   WDS_ARCH_NAMES = [nil, 'x86', 'ia64', 'x64', 'arm'].freeze
@@ -30,6 +32,7 @@ class ForemanWds::WdsImage
 
   def architecture_name
     return architecture unless architecture.is_a?(Integer) && WDS_ARCH_NAMES[architecture]
+
     WDS_ARCH_NAMES[architecture]
   end
 
@@ -43,6 +46,7 @@ class ForemanWds::WdsImage
 
   def matches_architecture?(architecture)
     return nil unless WDS_IMAGE_ARCHES[self.architecture]
+
     architecture = architecture.name if architecture.is_a? Architecture
     !(WDS_IMAGE_ARCHES[self.architecture] =~ architecture).nil?
   end
@@ -69,12 +73,11 @@ class ForemanWds::WdsImage
   end
 
   def parse_time(time)
-    return time unless time.is_a?(String) && time =~ /\/Date\(.*\)\//
+    return time unless time.is_a?(String) && time =~ %r{/Date\(.*\)/}
 
-    time = Time.at(time.scan(/\((.*)\)/).flatten.first.to_i / 1000)
+    time = Time.utc(time.scan(/\((.*)\)/).flatten.first.to_i / 1000.0)
     return time unless wds_server
 
-    time.utc
     sec_diff = wds_server.timezone - Time.at(0).utc_offset
     time += sec_diff
     time.localtime
diff --git a/app/lib/foreman_wds/wds_install_image.rb b/app/lib/foreman_wds/wds_install_image.rb
index 5866eb1d499d8d85114d5cbb86de3046b6866617..5a359f5b564b05ac2a9e7fdd3647f0ee50a792d9 100644
--- a/app/lib/foreman_wds/wds_install_image.rb
+++ b/app/lib/foreman_wds/wds_install_image.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class ForemanWds::WdsInstallImage < ForemanWds::WdsImage
   attr_accessor :compression, :dependent_files, :format, :image_group,
                 :partition_style, :security, :staged, :unattend_file_present
@@ -8,6 +10,7 @@ class ForemanWds::WdsInstallImage < ForemanWds::WdsImage
 
   def reload
     return false if wds_server.nil?
+
     @json = wds_server.install_image(name)
     load!
   end
diff --git a/app/models/concerns/foreman_wds/compute_resource_extensions.rb b/app/models/concerns/foreman_wds/compute_resource_extensions.rb
index b0467e6437b6c480add2b72929e27974e456549c..70583f46ee6e38715b286bbc75cc9373428d75eb 100644
--- a/app/models/concerns/foreman_wds/compute_resource_extensions.rb
+++ b/app/models/concerns/foreman_wds/compute_resource_extensions.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ForemanWds
   module ComputeResourceExtensions
     def capabilities
diff --git a/app/models/concerns/foreman_wds/host_extensions.rb b/app/models/concerns/foreman_wds/host_extensions.rb
index dcb216617dabc0d95f2222dacc7fd82982ea4f21..91bd2773e3d4843c63e1a0a56cbf1d06837e08a0 100644
--- a/app/models/concerns/foreman_wds/host_extensions.rb
+++ b/app/models/concerns/foreman_wds/host_extensions.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ForemanWds
   module HostExtensions
     def self.prepended(base)
diff --git a/app/models/concerns/foreman_wds/nic_extensions.rb b/app/models/concerns/foreman_wds/nic_extensions.rb
index 6781712d3903b55987a297c506d75bc85b512e39..3e34de37fa14e138346a50668e899007e2d8022a 100644
--- a/app/models/concerns/foreman_wds/nic_extensions.rb
+++ b/app/models/concerns/foreman_wds/nic_extensions.rb
@@ -1,10 +1,12 @@
+# frozen_string_literal: true
+
 module ForemanWds
   module NicExtensions
     def dhcp_update_required?
       return super if host.nil? || !host.wds? || host.wds_facet.nil?
 
       # DHCP entry for WDS depends on build mode
-      return true if host.build_changed?
+      true if host.build_changed?
     end
 
     def boot_server
@@ -12,6 +14,7 @@ module ForemanWds
 
       if host.build? # TODO: Support choosing local boot method
         return host.wds_server.next_server_ip unless subnet.dhcp.has_capability?(:DHCP, :dhcp_filename_hostname)
+
         return host.wds_server.next_server_name
       end
 
@@ -19,7 +22,8 @@ module ForemanWds
     end
 
     def dhcp_records
-      # Always recalculate dhcp records for WDS hosts, to allow different filename for the deleting and setting of a DHCP rebuild
+      # Always recalculate dhcp records for WDS hosts, to allow different filename
+      # for the deleting and setting of a DHCP rebuild
       @dhcp_records = nil if !host.nil? && host.wds?
       super
     end
diff --git a/app/models/foreman_wds/wds_facet.rb b/app/models/foreman_wds/wds_facet.rb
index bbc836021eb07c3f798e73f9dc91e4b8eebf1c6c..a117887f739be750fbe1c4670ebda174b7d69c6c 100644
--- a/app/models/foreman_wds/wds_facet.rb
+++ b/app/models/foreman_wds/wds_facet.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ForemanWds
   class WdsFacet < ApplicationRecord
     class Jail < Safemode::Jail
@@ -13,11 +15,11 @@ module ForemanWds
     validates_lengths_from_database
 
     validates :host, presence: true, allow_blank: false
-
     validates :install_image_name, presence: true, allow_blank: false
 
     def boot_image
       return nil unless wds_server
+
       @boot_image ||= if boot_image_name
                         wds_server.boot_image(boot_image_name)
                       else
@@ -27,21 +29,25 @@ module ForemanWds
 
     def boot_image_file
       return nil unless boot_image
+
       @boot_image_file ||= boot_image.file_name
     end
 
     def install_image
       return nil unless wds_server
+
       @install_image ||= wds_server.install_image(install_image_name)
     end
 
     def install_image_file
       return nil unless install_image
+
       @install_image_file ||= install_image.file_name
     end
 
     def install_image_group
       return nil unless install_image
+
       @install_image_group ||= install_image.image_group
     end
   end
diff --git a/app/models/wds_server.rb b/app/models/wds_server.rb
index 36eddf9eadd0201d3e232ab3ac91a8e0352a7880..aaf1cb4c301a4197211f2dfe210dbe53c182f21e 100644
--- a/app/models/wds_server.rb
+++ b/app/models/wds_server.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class WdsServer < ApplicationRecord
   class Jail < Safemode::Jail
     allow :name, :shortname
@@ -44,8 +46,8 @@ class WdsServer < ApplicationRecord
   end
 
   def clients
-    objects = connection.run_wql('SELECT * FROM MSFT_WdsClient')[:msft_wdsclient] rescue nil
-    objects = nil if objects.empty?
+    objects = run_wql('SELECT * FROM MSFT_WdsClient', on_error: {})[:msft_wdsclient]
+    objects = nil if objects&.empty?
     objects ||= begin
       data = connection.shell(:powershell) do |s|
         s.run('Get-WdsClient | ConvertTo-Json -Compress')
@@ -66,14 +68,14 @@ class WdsServer < ApplicationRecord
 
   def boot_images
     cache.cache(:boot_images) do
-      images(:boot)
-    end.each { |i| i.wds_server = self }
+      images(:boot).each { |i| i.wds_server = self }
+    end
   end
 
   def install_images
     cache.cache(:install_images) do
-      images(:install)
-    end.each { |i| i.wds_server = self }
+      images(:install).each { |i| i.wds_server = self }
+    end
   end
 
   def create_client(host)
@@ -100,13 +102,13 @@ class WdsServer < ApplicationRecord
 
   def timezone
     cache.cache(:timezone) do
-      connection.run_wql('SELECT Bias FROM Win32_TimeZone')[:xml_fragment].first[:bias].to_i * 60
+      run_wql('SELECT Bias FROM Win32_TimeZone')[:xml_fragment].first[:bias].to_i * 60
     end
   end
 
   def shortname
     cache.cache(:shortname) do
-      connection.run_wql('SELECT Name FROM Win32_ComputerSystem')[:xml_fragment].first[:name]
+      run_wql('SELECT Name FROM Win32_ComputerSystem')[:xml_fragment].first[:name]
     end
   end
 
@@ -149,9 +151,7 @@ class WdsServer < ApplicationRecord
   end
 
   def test_connection
-    connection.run_wql('SELECT * FROM Win32_UTCTime').key? :win32_utc_time
-  rescue StandardError
-    false
+    run_wql('SELECT * FROM Win32_UTCTime', on_error: {}).key? :win32_utc_time
   end
 
   def refresh_cache
@@ -185,6 +185,7 @@ class WdsServer < ApplicationRecord
   def ensure_unattend(host)
     iface = host&.provision_interface
     raise 'No provisioning interface available' unless iface
+
     raise NotImplementedException, 'TODO: Not implemented yet'
 
     # TODO: render template, send as heredoc
@@ -226,23 +227,19 @@ class WdsServer < ApplicationRecord
     end.errcode.zero?
   end
 
-  def ensure_client(host)
+  def ensure_client(_host)
     raise NotImplementedException, 'TODO: Not implemented yet'
   end
 
-  def delete_client(host)
+  def delete_client(_host)
     raise NotImplementedException, 'TODO: Not implemented yet'
   end
 
   def images(type, name = nil)
     raise ArgumentError, 'Type must be :boot or :install' unless %i[boot install].include? type
 
-    begin
-      objects = connection.run_wql("SELECT * FROM MSFT_Wds#{type.to_s.capitalize}Image#{" WHERE Name=\"#{name}\"" if name}")["msft_wds#{type}image".to_sym]
-      objects = nil if objects.empty?
-    rescue StandardError => e
-      ::Rails.logger.debug "WQL image query failed with #{e.class}: #{e}"
-    end
+    objects = run_wql("SELECT * FROM MSFT_Wds#{type.to_s.capitalize}Image#{" WHERE Name=\"#{name}\"" if name}", on_error: {})["msft_wds#{type}image".to_sym]
+    objects = nil if objects.empty?
 
     unless objects
       begin
@@ -267,12 +264,20 @@ class WdsServer < ApplicationRecord
     when Array
       result.map { |v| underscore_result(v) }
     when Hash
-      Hash[result.map { |k, v| [k.to_s.underscore.to_sym, underscore_result(v)] }]
+      result.to_h { |k, v| [k.to_s.underscore.to_sym, underscore_result(v)] }
     else
       result
     end
   end
 
+  def run_wql(wql, on_error: :raise)
+    connection.run_wql(wql)
+  rescue StandardError
+    raise if on_error == :raise
+
+    on_error
+  end
+
   def cache
     @cache ||= WdsImageCache.new(self)
   end
diff --git a/app/services/wds_image_cache.rb b/app/services/wds_image_cache.rb
index 25fd510ef3aa1d5666c924f44b4ff8475c2c6967..901436925badafa5f874bc5d563ef1e4b852b1ad 100644
--- a/app/services/wds_image_cache.rb
+++ b/app/services/wds_image_cache.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 class WdsImageCache
   attr_accessor :wds_server, :cache_duration
 
@@ -46,6 +48,7 @@ class WdsImageCache
 
   def get_uncached_value(&block)
     return unless block_given?
+
     wds_server.instance_eval(&block)
   end
 
diff --git a/config/routes.rb b/config/routes.rb
index ec6bcf6e01a5de5ed955198d1a3b39738da99de0..bbcd9358ece8ee6bb4b02714b9e4da3b112b4a08 100644
--- a/config/routes.rb
+++ b/config/routes.rb
@@ -1,4 +1,6 @@
 Rails.application.routes.draw do
+  post 'wds_server_selected_discovered_hosts' => 'hosts#wds_server_selected'
+
   scope '/foreman_wds' do
     constraints(id: %r{[^\/]+}) do
       resources :wds_servers do
diff --git a/lib/foreman_wds.rb b/lib/foreman_wds.rb
index d992bd7c16eba364874f49c15912e07ba6c16cfc..79ed42b2bba9d3d18f20f1462d9770c7c264cf19 100644
--- a/lib/foreman_wds.rb
+++ b/lib/foreman_wds.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'foreman_wds/engine'
 
 module ForemanWds
diff --git a/lib/foreman_wds/engine.rb b/lib/foreman_wds/engine.rb
index be7ac158d5cdcd8f1f533caaa54ca636724c65d4..0dd95334eb86928c1ee2cc1d4042407a9872ee7f 100644
--- a/lib/foreman_wds/engine.rb
+++ b/lib/foreman_wds/engine.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ForemanWds
   class Engine < ::Rails::Engine
     engine_name 'foreman_wds'
@@ -36,9 +38,9 @@ module ForemanWds
           # }, resource_type: 'Host'
         end
 
-        Foreman::AccessControl.permission(:edit_hosts).actions.concat [
+        Foreman::AccessControl.permission(:edit_hosts).actions.push(
           'hosts/wds_server_selected', 'hosts/wds_image_selected'
-        ]
+        )
 
         role 'WDS Server Manager',
              %i[view_wds_servers create_wds_servers edit_wds_servers destroy_wds_servers],
@@ -76,18 +78,16 @@ module ForemanWds
     end
 
     config.to_prepare do
-      begin
-        Host::Managed.send(:prepend, ForemanWds::HostExtensions)
-        Nic::Managed.send(:prepend, ForemanWds::NicExtensions)
-        HostsController.send(:prepend, ForemanWds::HostsControllerExtensions)
-        UnattendedController.send(:prepend, ForemanWds::UnattendedControllerExtensions)
-
-        ComputeResource.providers.each do |_k, const|
-          Kernel.const_get(const).send(:prepend, ForemanWds::ComputeResourceExtensions)
-        end
-      rescue StandardError => e
-        Rails.logger.fatal "foreman_wds: skipping engine hook (#{e})"
+      Host::Managed.prepend ForemanWds::HostExtensions
+      Nic::Managed.prepend ForemanWds::NicExtensions
+      HostsController.include ForemanWds::HostsControllerExtensions
+      UnattendedController.prepend ForemanWds::UnattendedControllerExtensions
+
+      ComputeResource.providers.each do |_k, const|
+        Kernel.const_get(const).send(:prepend, ForemanWds::ComputeResourceExtensions)
       end
+    rescue StandardError => e
+      Rails.logger.fatal "foreman_wds: skipping engine hook (#{e})"
     end
 
     rake_tasks do
diff --git a/lib/tasks/foreman_wds_tasks.rake b/lib/tasks/foreman_wds_tasks.rake
new file mode 100644
index 0000000000000000000000000000000000000000..13524301426751b3acd0a605e22eef7029fd19c4
--- /dev/null
+++ b/lib/tasks/foreman_wds_tasks.rake
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require 'rake/testtask'
+
+# Tests
+namespace :test do
+  desc "Test ForemanWDS"
+  Rake::TestTask.new(:foreman_wds) do |t|
+    test_dir = File.join(File.dirname(File.dirname(__dir__)), 'test')
+    t.libs << ['test', test_dir]
+    t.pattern = "#{test_dir}/**/*_test.rb"
+    t.verbose = true
+    t.warning = false
+  end
+
+  namespace :foreman_wds do
+    task :coverage do
+      ENV['COVERAGE'] = '1'
+
+      Rake::Task['test:foreman_wds'].invoke
+    end
+  end
+end
+
+namespace :foreman_wds do
+  task rubocop: :environment do
+    begin
+      require 'rubocop/rake_task'
+      RuboCop::RakeTask.new(:rubocop_foreman_wds) do |task|
+        task.patterns = ["#{ForemanWDS::Engine.root}/app/**/*.rb",
+                         "#{ForemanWDS::Engine.root}/lib/**/*.rb"]
+      end
+    rescue StandardError
+      puts 'Rubocop not loaded.'
+    end
+
+    Rake::Task['rubocop_foreman_wds'].invoke
+  end
+end
+
+Rake::Task[:test].enhance ['test:foreman_wds']
+# Rake::Task[:coverage].enhance ["test:foreman_liudesk_cmdb:coverage"]
diff --git a/test/factories/host.rb b/test/factories/host.rb
new file mode 100644
index 0000000000000000000000000000000000000000..c15dec6792871ec35e2be9563ae477810008ca7b
--- /dev/null
+++ b/test/factories/host.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+FactoryBot.modify do
+  factory :host do
+    trait :with_wds_facet do
+      association :wds_facet, factory: :wds_facet, strategy: :build
+    end
+  end
+end
diff --git a/test/factories/wds_facet.rb b/test/factories/wds_facet.rb
new file mode 100644
index 0000000000000000000000000000000000000000..c8b00afa32b2ddd0f268fa52b4a87e4ecec4a7f8
--- /dev/null
+++ b/test/factories/wds_facet.rb
@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+FactoryBot.define do
+  factory :wds_facet, class: 'ForemanWds::WdsFacet' do
+    host
+    wds_server
+  end
+end
diff --git a/test/factories/wds_server.rb b/test/factories/wds_server.rb
new file mode 100644
index 0000000000000000000000000000000000000000..9b80ec7a50b4914180bc198a61a04316b32c7764
--- /dev/null
+++ b/test/factories/wds_server.rb
@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+FactoryBot.define do
+  factory :wds_server, class: 'WdsServer' do
+    name { 'WDS server' }
+    description { 'Example WDS server for testing' }
+    url { 'http://example.com:5985/wsman' }
+    user { 'username' }
+    password { 'password' }
+  end
+end
diff --git a/test/foreman_wds_test.rb b/test/foreman_wds_test.rb
deleted file mode 100644
index e89cc158bca69c6a82e6edb00997898ee73b20e0..0000000000000000000000000000000000000000
--- a/test/foreman_wds_test.rb
+++ /dev/null
@@ -1,11 +0,0 @@
-require "test_helper"
-
-class ForemanWdsTest < Minitest::Test
-  def test_that_it_has_a_version_number
-    refute_nil ::ForemanWds::VERSION
-  end
-
-  def test_it_does_something_useful
-    assert false
-  end
-end
diff --git a/test/test_helper.rb b/test/test_helper.rb
deleted file mode 100644
index 5b94c7573f2d7eda759c218490d332d3756be2a7..0000000000000000000000000000000000000000
--- a/test/test_helper.rb
+++ /dev/null
@@ -1,4 +0,0 @@
-$LOAD_PATH.unshift File.expand_path("../../lib", __FILE__)
-require "foreman_wds"
-
-require "minitest/autorun"
diff --git a/test/test_plugin_helper.rb b/test/test_plugin_helper.rb
new file mode 100644
index 0000000000000000000000000000000000000000..a236a6dbe659d84b495b26b40d599aefcd81975d
--- /dev/null
+++ b/test/test_plugin_helper.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+if ENV['COVERAGE']
+  require 'simplecov'
+
+  SimpleCov.start 'rails' do
+    root File.dirname(__dir__)
+
+    add_group 'Interactors', '/app/interactors'
+    add_group 'Services', '/app/services'
+
+    formatter SimpleCov::Formatter::SimpleFormatter if ENV['CI']
+  end
+end
+
+# This calls the main test_helper in Foreman-core
+require 'test_helper'
+
+ActiveSupport::TestCase.file_fixture_path = File.join(__dir__, 'fixtures')
+
+# Add plugin to FactoryBot's paths
+FactoryBot.definition_file_paths << File.join(__dir__, 'factories')
+FactoryBot.reload