(PUP-3244) Validate Environment Before Setting

Before this commit the environment validation was happening for the Exec
terminus instead of on the node object itself. This fix should help all
node objects from having invalid environments when its environment is
set using `Puppet::Node#environment=`.

2 files changed, 6 insertions, 9 deletions

diff --git a/lib/puppet/indirector/node/exec.rb b/lib/puppet/indirector/node/exec.rb
index 9fa3c07ad..d4faf74f4 100644
--- a/lib/puppet/indirector/node/exec.rb
+++ b/lib/puppet/indirector/node/exec.rb
@@ -19,14 +19,6 @@ class Puppet::Node::Exec < Puppet::Indirector::Exec
     # Translate the output to ruby.
     result = translate(request.key, output)
 
-    # If defining a new environment then we check to see if
-    # the environment exists, if it does not then we raise
-    # an error instead of falling back to the 'production'
-    # environment.
-    if result[:environment] && !Puppet.lookup(:environments).get(result[:environment])
-      raise Puppet::Environments::EnvironmentNotFound, result[:environment]
-    end
-
     # Set the requested environment if it wasn't overridden
     # If we don't do this it gets set to the local default
     result[:environment] ||= request.environment.name
diff --git a/lib/puppet/node.rb b/lib/puppet/node.rb
index d61d49385..8f79224b1 100644
--- a/lib/puppet/node.rb
+++ b/lib/puppet/node.rb
@@ -76,9 +76,14 @@ 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)
-    else
+    elsif env.is_a?(Puppet::Node::Environment)
       @environment = env
+    else
+        raise Puppet::Environments::EnvironmentNotFound, env
     end
   end