3 files changed, 39 insertions, 0 deletions
diff --git a/manifests/params.pp b/manifests/params.pp
index 3d0b0f5..cce558d 100644
--- a/manifests/params.pp
+++ b/manifests/params.pp
@@ -6,6 +6,9 @@ class augeas::params {
 
   if versioncmp($::puppetversion, '4.0.0') >= 0 {
     $lens_dir = '/opt/puppetlabs/puppet/share/augeas/lenses'
+  } elsif (str2bool($::is_pe)) {
+    # puppet enterpise has a different lens location
+    $lens_dir = '/opt/puppet/share/augeas/lenses'
   } else {
     $lens_dir = '/usr/share/augeas/lenses'
   }
diff --git a/spec/classes/augeas_spec.rb b/spec/classes/augeas_spec.rb
index 32a390d..b7c8273 100644
--- a/spec/classes/augeas_spec.rb
+++ b/spec/classes/augeas_spec.rb
@@ -7,6 +7,7 @@ describe 'augeas' do
       {
         :osfamily      => 'MS-DOS',
         :puppetversion => Puppet.version,
+        :is_pe         => false,
       }
     end
 
@@ -22,6 +23,7 @@ describe 'augeas' do
       let(:facts) do
         facts.merge({
           :puppetversion => Puppet.version,
+          :is_pe         => false,
         })
       end
 
@@ -152,6 +154,39 @@ describe 'augeas' do
           :force        => 'true'
         ).without(:recurse) }
       end
+
+      context 'with Puppet Enterprise' do
+        let (:facts) do
+          facts.merge({
+            :puppetversion => Puppet.version,
+            :is_pe         => true,
+          })
+        end
+
+        if Puppet::Util::Package.versioncmp(Puppet.version, '4.0.0') >= 0
+            # the enterprise lens dir is the same in 4
+            pe_lens_dir = lens_dir
+        else
+            pe_lens_dir = '/opt/puppet/share/augeas/lenses'
+        end
+
+        it { is_expected.to contain_file(pe_lens_dir).with(
+          :ensure       => 'directory',
+          :force        => 'true',
+          :recurse      => 'true',
+          :recurselimit => 1
+        ) }
+        it { is_expected.to contain_file("#{pe_lens_dir}/dist").with(
+          :ensure  => 'directory',
+          :purge   => 'false'
+        ) }
+        it { is_expected.to contain_file("#{pe_lens_dir}/tests").with(
+          :ensure => 'directory',
+          :force  => 'true',
+          :purge  => 'true'
+        ).without(:recurse) }
+      end
+
     end
   end
 end
diff --git a/spec/defines/augeas_lens_spec.rb b/spec/defines/augeas_lens_spec.rb
index f4e5be7..57747ca 100644
--- a/spec/defines/augeas_lens_spec.rb
+++ b/spec/defines/augeas_lens_spec.rb
@@ -25,6 +25,7 @@ describe 'augeas::lens' do
           facts.merge({
             :augeasversion => :undef,
             :puppetversion => Puppet.version,
+            :is_pe         => false,
           })
         end