(PUP-3244) Add Puppet::Environment::Directories#get! method

This new function will return the directory environment or raise an
error if a directory environment does not exist. This gets rid of the
need to check if environments exist across the codebase before
continuing for the areas where we need.

3 files changed, 67 insertions, 7 deletions

diff --git a/lib/puppet/environments.rb b/lib/puppet/environments.rb
index 9f7ac5c31..6af5622de 100644
--- a/lib/puppet/environments.rb
+++ b/lib/puppet/environments.rb
@@ -48,6 +48,14 @@ module Puppet::Environments
   #     we are looking up
   #   @return [Puppet::Setting::EnvironmentConf, nil] the configuration for the
   #     requested environment, or nil if not found or no configuration is available
+  #
+  # @!macro [new] loader_get_or_fail
+  #   Find a named environment or raise
+  #   Puppet::Environments::EnvironmentNotFound when the named environment is
+  #   does not exist.
+  #
+  #   @param name [String,Symbol] The name of environment to find
+  #   @return [Puppet::Node::Environment] the requested environment
 
   # A source of pre-defined environments.
   #
@@ -76,6 +84,14 @@ module Puppet::Environments
       end
     end
 
+    # @!macro loader_get_or_fail
+    def get!(name)
+      if !environment = get(name)
+        raise EnvironmentNotFound, name
+      end
+      environment
+    end
+
     # Returns a basic environment configuration object tied to the environment's
     # implementation values.  Will not interpolate.
     #
@@ -144,6 +160,14 @@ module Puppet::Environments
       Puppet::Node::Environment.new(name)
     end
 
+    # @note Because the Legacy system cannot list out all of its environments,
+    #   this method will never fail and is only calling get directly.
+    #
+    # @!macro loader_get_or_fail
+    def get!(name)
+      get(name)
+    end
+
     # @note we could return something here, but since legacy environments
     #   are deprecated, there is no point.
     #
@@ -206,6 +230,14 @@ module Puppet::Environments
       list.find { |env| env.name == name.intern }
     end
 
+    # @!macro loader_get_or_fail
+    def get!(name)
+      if !environment = get(name)
+        raise EnvironmentNotFound, name
+      end
+      environment
+    end
+
     # @!macro loader_get_conf
     def get_conf(name)
       valid_directories.each do |envdir|
@@ -259,6 +291,14 @@ module Puppet::Environments
       nil
     end
 
+    # @!macro loader_get_or_fail
+    def get!(name)
+      if !environment = get(name)
+        raise EnvironmentNotFound, name
+      end
+      environment
+    end
+
     # @!macro loader_get_conf
     def get_conf(name)
       @loaders.each do |loader|
@@ -289,6 +329,14 @@ module Puppet::Environments
       end
     end
 
+    # @!macro loader_get_or_fail
+    def get!(name)
+      if !environment = get(name)
+        raise EnvironmentNotFound, name
+      end
+      environment
+    end
+
     # Clears the cache of the environment with the given name.
     # (The intention is that this could be used from a MANUAL cache eviction command (TBD)
     def clear(name)
diff --git a/lib/puppet/node.rb b/lib/puppet/node.rb
index 8f79224b1..6f0d89dba 100644
--- a/lib/puppet/node.rb
+++ b/lib/puppet/node.rb
@@ -76,14 +76,9 @@ class Puppet::Node
 
   def environment=(env)
     if env.is_a?(String) or env.is_a?(Symbol)
-      if !Puppet.lookup(:environments).get(env)
-        raise Puppet::Environments::EnvironmentNotFound, env
-      end
-      @environment = Puppet.lookup(:environments).get(env)
-    elsif env.is_a?(Puppet::Node::Environment)
-      @environment = env
+      @environment = Puppet.lookup(:environments).get!(env)
     else
-        raise Puppet::Environments::EnvironmentNotFound, env
+      @environment = env
     end
   end
 
diff --git a/spec/unit/environments_spec.rb b/spec/unit/environments_spec.rb
index 80a0bb2d5..917790f29 100644
--- a/spec/unit/environments_spec.rb
+++ b/spec/unit/environments_spec.rb
@@ -106,6 +106,17 @@ describe Puppet::Environments do
       end
     end
 
+    it "raises error if an environment can't be found" do
+      directory_tree = FS::MemoryFile.a_directory("envdir", [])
+
+      loader_from(:filesystem => [directory_tree],
+                  :directory => directory_tree) do |loader|
+        expect do
+          loader.get!("env_not_in_this_list")
+        end.to raise_error(Puppet::Environments::EnvironmentNotFound)
+      end
+    end
+
     context "with an environment.conf" do
       let(:envdir) do
         FS::MemoryFile.a_directory(File.expand_path("envdir"), [
@@ -315,6 +326,12 @@ config_version=$vardir/random/scripts
       expect(loader.get(:doesnotexist)).to be_nil
     end
 
+    it "raises error if environment is not found" do
+      expect do
+        loader.get!(:doesnotexist)
+      end.to raise_error(Puppet::Environments::EnvironmentNotFound)
+    end
+
     it "gets a basic conf" do
       conf = loader.get_conf(:static1)
       expect(conf.modulepath).to eq('')