diff options
| -rw-r--r-- | lib/puppet/node/environment.rb | 14 | ||||
| -rwxr-xr-x | spec/unit/node/environment.rb | 24 |
2 files changed, 38 insertions, 0 deletions
diff --git a/lib/puppet/node/environment.rb b/lib/puppet/node/environment.rb index 537a5dd26..70417449e 100644 --- a/lib/puppet/node/environment.rb +++ b/lib/puppet/node/environment.rb @@ -7,6 +7,20 @@ end # Model the environment that a node can operate in. This class just # provides a simple wrapper for the functionality around environments. class Puppet::Node::Environment + module Helper + def environment + Puppet::Node::Environment.new(@environment) + end + + def environment=(env) + if env.is_a?(String) or env.is_a?(Symbol) + @environment = env + else + @environment = env.name + end + end + end + include Puppet::Util::Cacher @seen = {} diff --git a/spec/unit/node/environment.rb b/spec/unit/node/environment.rb index 81981973d..8945f28df 100755 --- a/spec/unit/node/environment.rb +++ b/spec/unit/node/environment.rb @@ -216,4 +216,28 @@ describe Puppet::Node::Environment do end end end + + describe Puppet::Node::Environment::Helper do + before do + @helper = Object.new + @helper.extend(Puppet::Node::Environment::Helper) + end + + it "should be able to set and retrieve the environment" do + @helper.environment = :foo + @helper.environment.name.should == :foo + end + + it "should accept an environment directly" do + env = Puppet::Node::Environment.new :foo + @helper.environment = env + @helper.environment.name.should == :foo + end + + it "should accept an environment as a string" do + env = Puppet::Node::Environment.new "foo" + @helper.environment = env + @helper.environment.name.should == :foo + end + end end |