Fixing the SSL::Host#waitforcert method.

It now works the way puppetd needs it to, rather
than the way I thought it would need to work.

Signed-off-by: Luke Kanies <luke@madstop.com>

2 files changed, 63 insertions, 23 deletions

diff --git a/lib/puppet/ssl/host.rb b/lib/puppet/ssl/host.rb
index a449dcc7e..d3805eb20 100644
--- a/lib/puppet/ssl/host.rb
+++ b/lib/puppet/ssl/host.rb
@@ -187,18 +187,34 @@ class Puppet::SSL::Host
     # Attempt to retrieve a cert, if we don't already have one.
     def wait_for_cert(time)
         return :existing if certificate
-        exit(1) if time < 1
-        generate_certificate_request
+        begin
+            generate
+
+            return :new if certificate
+        rescue StandardError => detail
+            Puppet.err "Could not request certificate: %s" % detail.to_s
+            if time < 1
+                puts "Exiting; failed to retrieve certificate and watiforcert is disabled"
+                exit(1)
+            else
+                sleep(time)
+            end
+            retry
+        end
+
+        if time < 1
+            puts "Exiting; no certificate found and waitforcert is disabled"
+            exit(1) 
+        end
 
         while true do
-           begin
-               break if certificate
-               Puppet.notice "Did not receive certificate"
-           rescue StandardError => detail
-               Puppet.err "Could not request certificate: %s" % detail.to_s
-           end
-
-           sleep time
+            sleep time
+            begin
+                break if certificate
+                Puppet.notice "Did not receive certificate"
+            rescue StandardError => detail
+                Puppet.err "Could not request certificate: %s" % detail.to_s
+            end
         end
         return :new
     end
diff --git a/spec/unit/ssl/host.rb b/spec/unit/ssl/host.rb
index 73ce9c7f2..c234585f7 100755
--- a/spec/unit/ssl/host.rb
+++ b/spec/unit/ssl/host.rb
@@ -455,36 +455,60 @@ describe Puppet::SSL::Host do
             @host.expects(:certificate).returns "foo"
             @host.wait_for_cert(0).should == :existing
         end
-        
-        it "should exit if it has no certificate and the wait time is 0" do
-            @host.expects(:certificate).returns nil
-            @host.expects(:exit).with(1).raises(SystemExit)
-            lambda { @host.wait_for_cert(0) }.should raise_error(SystemExit)
-        end
 
         it "should generate its certificate request and attempt to read the certificate again if no certificate is found" do
             @host.expects(:certificate).times(2).returns(nil).then.returns "foo"
-            @host.expects(:generate_certificate_request)
-            @host.wait_for_cert(10).should == :new
+            @host.expects(:generate)
+            @host.wait_for_cert(1).should == :new
+        end
+
+        it "should catch and log errors during CSR saving" do
+            @host.expects(:certificate).times(2).returns(nil).then.returns "foo"
+            @host.expects(:generate).times(2).raises(RuntimeError).then.returns nil
+            @host.stubs(:sleep)
+            @host.wait_for_cert(1).should == :new
+        end
+
+        it "should sleep and retry after failures saving the CSR if waitforcert is enabled" do
+            @host.expects(:certificate).times(2).returns(nil).then.returns "foo"
+            @host.expects(:generate).times(2).raises(RuntimeError).then.returns nil
+            @host.expects(:sleep).with(1)
+            @host.wait_for_cert(1)
+        end
+
+        it "should exit after failures saving the CSR of waitforcert is disabled" do
+            @host.expects(:certificate).returns(nil)
+            @host.expects(:generate).raises(RuntimeError)
+            @host.expects(:puts)
+            @host.expects(:exit).with(1).raises(SystemExit)
+            lambda { @host.wait_for_cert(0) }.should raise_error(SystemExit)
+        end
+        
+        it "should exit if the wait time is 0 and it can neither find nor retrieve a certificate" do
+            @host.stubs(:certificate).returns nil
+            @host.expects(:generate)
+            @host.expects(:puts)
+            @host.expects(:exit).with(1).raises(SystemExit)
+            lambda { @host.wait_for_cert(0) }.should raise_error(SystemExit)
         end
 
         it "should sleep for the specified amount of time if no certificate is found after generating its certificate request" do
             @host.expects(:certificate).times(3).returns(nil).then.returns(nil).then.returns "foo"
-            @host.expects(:generate_certificate_request)
+            @host.expects(:generate)
 
-            @host.expects(:sleep).with(10)
+            @host.expects(:sleep).with(1)
 
-            @host.wait_for_cert(10).should == :new
+            @host.wait_for_cert(1).should == :new
         end
 
         it "should catch and log exceptions during certificate retrieval" do
             @host.expects(:certificate).times(3).returns(nil).then.raises(RuntimeError).then.returns("foo")
-            @host.stubs(:generate_certificate_request)
+            @host.stubs(:generate)
             @host.stubs(:sleep)
 
             Puppet.expects(:err)
 
-            @host.wait_for_cert(10)
+            @host.wait_for_cert(1)
         end
     end
 end