Finalizing the filebucket client, with test code.

git-svn-id: https://reductivelabs.com/svn/puppet/trunk@2348 980ebf18-57e1-0310-9a29-db15c13687c0

bin/pbucket

@@ -16,6 +16,36 @@
 # This is a stand-alone filebucket client for sending files to a local
 # or central filebucket.
 #
+# = Usage
+#
+# This client can operate in three modes, with only one mode per call:
+#
+# backup::
+#   Send one or more files to the specified file bucket.  Each sent file
+#   is printed with its resulting md5 sum.
+#
+# get::
+#   Return the text associated with an md5 sum.  The text is printed to
+#   stdout, and only one file can be retrieved at a time.
+#
+# restore::
+#   Given a file path and an md5 sum, store the content associated with the
+#   sum into the specified file path.  You can specify an entirely new path
+#   to this argument; you are not restricted to restoring the content to its
+#   original location.
+#
+# Note that +pbucket+ defaults to using a network-based filebucket available on
+# the server named +puppet+.  To use this, you'll have to be running as a user
+# with valid Puppet certificates.  Alternatively, you can use your local file bucket
+# by specifying +--local+.
+#
+# = Example
+#
+#   $ pbucket backup /etc/passwd
+#   /etc/passwd: 429b225650b912a2ee067b0a4cf1e949
+#   $ pbucket restore /tmp/passwd 429b225650b912a2ee067b0a4cf1e949
+#   $
+# 
 # = Options
 #
 # Note that any configuration parameter that's valid in the configuration file
@@ -65,8 +95,7 @@
 # Licensed under the GNU Public License
 
 require 'puppet'
-require 'puppet/server'
-require 'puppet/client'
+require 'puppet/network/client'
 require 'getoptlong'
 
 options = [
@@ -140,11 +169,12 @@ Puppet.genconfig
 Puppet.genmanifest
 
 begin
-    if options[:local]
+    if options[:local] or options[:bucket]
         path = options[:bucket] || Puppet[:bucketdir]
-        client = Puppet::Client::Dipper.new(:Path => path)
+        client = Puppet::Network::Client.dipper.new(:Path => path)
     else
-        client = Puppet::Client::Dipper.new(:Server => Puppet[:server])
+        require 'puppet/network/handler'
+        client = Puppet::Network::Client.dipper.new(:Server => Puppet[:server])
     end
 rescue => detail
     $stderr.puts detail
@@ -157,7 +187,10 @@ end
 mode = ARGV.shift
 case mode
 when "get":
-when "send":
+    md5 = ARGV.shift
+    out = client.getfile(md5)
+    print out
+when "backup":
     ARGV.each do |file|
         unless FileTest.exists?(file)
             $stderr.puts "%s: no such file" % file
@@ -167,8 +200,13 @@ when "send":
             $stderr.puts "%s: cannot read file" % file
             next
         end
-        client.backup(file)
+        md5 = client.backup(file)
+        puts "%s: %s" % [file, md5]
     end
+when "restore":
+    file = ARGV.shift
+    md5 = ARGV.shift
+    client.restore(file, md5)
 else
     raise "Invalid mode %s" % mode
 end

lib/puppet/network/client/dipper.rb

@@ -28,6 +28,17 @@ class Puppet::Network::Client::Dipper < Puppet::Network::Client
         return @driver.addfile(contents,file)
     end
 
+    # Retrieve a file by sum.
+    def getfile(sum)
+        if newcontents = @driver.getfile(sum)
+            unless local?
+                newcontents = Base64.decode64(newcontents)
+            end
+            return newcontents
+        end
+        return nil
+    end
+
     # Restore the file
     def restore(file,sum)
         restore = true
@@ -42,18 +53,15 @@ class Puppet::Network::Client::Dipper < Puppet::Network::Client
         end
 
         if restore
-            if newcontents = @driver.getfile(sum)
-                unless local?
-                    newcontents = Base64.decode64(newcontents)
-                end
+            if newcontents = getfile(sum)
                 tmp = ""
                 newsum = Digest::MD5.hexdigest(newcontents)
                 changed = nil
-                unless FileTest.writable?(file)
+                if FileTest.exists?(file) and ! FileTest.writable?(file)
                     changed = ::File.stat(file).mode
                     ::File.chmod(changed | 0200, file)
                 end
-                ::File.open(file,::File::WRONLY|::File::TRUNC) { |of|
+                ::File.open(file, ::File::WRONLY|::File::TRUNC|::File::CREAT) { |of|
                     of.print(newcontents)
                 }
                 if changed

lib/puppet/network/handler/filebucket.rb

@@ -28,7 +28,7 @@ class Puppet::Network::Handler # :nodoc:
         # this doesn't work for relative paths
         def self.paths(base,md5)
             dir = File.join(md5[0..7].split(""))
-            basedir = File.join(base, dir)
+            basedir = File.join(base, dir, md5)
             return [
                 basedir,
                 File.join(basedir, "contents"),

test/executables/filebucket.rb

@@ -0,0 +1,49 @@
+#!/usr/bin/env ruby
+
+$:.unshift("../lib").unshift("../../lib") if __FILE__ =~ /\.rb$/
+
+require 'puppet'
+require 'puppet/network/client'
+require 'puppettest'
+require 'socket'
+require 'facter'
+
+class TestPBucket < Test::Unit::TestCase
+    include PuppetTest::ExeTest
+
+    def test_local
+        bucket = tempfile
+        file = tempfile
+        text = "somet ext"
+        md5 = Digest::MD5.hexdigest(text)
+        File.open(file, "w") { |f| f.print text }
+        out = %x{pbucket --bucket #{bucket} backup #{file}}
+
+        outfile, outmd5 = out.chomp.split(": ")
+
+        assert_equal(0, $?, "pbucket did not run successfully")
+
+        assert_equal(file, outfile, "did not output correct file name")
+        assert_equal(md5, outmd5, "did not output correct md5 sum")
+
+        dipper = Puppet::Network::Client.dipper.new(:Path => bucket)
+
+        newtext = nil
+        assert_nothing_raised("Could not get file from bucket") do
+            newtext = dipper.getfile(md5)
+        end
+
+        assert_equal(text, newtext, "did not get correct file from md5 sum")
+
+        out = %x{pbucket --bucket #{bucket} get #{md5}}
+        assert_equal(0, $?, "pbucket did not run successfully")
+        assert_equal(text, out, "did not get correct text back from pbucket")
+
+        File.open(file, "w") { |f| f.puts "some other txt" }
+        out = %x{pbucket --bucket #{bucket} restore #{file} #{md5}}
+        assert_equal(0, $?, "pbucket did not run successfully")
+        assert_equal(text, File.read(file), "file was not restored")
+    end
+end
+
+# $Id$

test/network/client/dipper.rb

@@ -0,0 +1,35 @@
+#!/usr/bin/env ruby
+
+$:.unshift("../../lib") if __FILE__ =~ /\.rb$/
+
+require 'puppettest'
+require 'puppet/network/client/resource'
+
+class TestDipperClient < Test::Unit::TestCase
+    include PuppetTest::ServerTest
+
+    def setup
+        super
+        @dipper = Puppet::Network::Client.dipper.new(:Path => tempfile)
+    end
+
+    # Make sure we can create a new file with 'restore'.
+    def test_restore_to_new_file
+        file = tempfile
+        text = "asdf;lkajseofiqwekj"
+        File.open(file, "w") { |f| f.puts text }
+        md5 = nil
+        assert_nothing_raised("Could not send file") do
+            md5 = @dipper.backup(file)
+        end
+
+        newfile = tempfile
+        assert_nothing_raised("could not restore to new path") do
+            @dipper.restore(newfile, md5)
+        end
+
+        assert_equal(File.read(file), File.read(newfile), "did not restore correctly")
+    end
+end
+
+# $Id$

test/network/handler/bucket.rb

@@ -276,7 +276,7 @@ class TestBucket < Test::Unit::TestCase
         md5 = Digest::MD5.hexdigest(text)
 
         dirs = File.join(md5[0..7].split(""))
-        dir = File.join(@bucket, dirs)
+        dir = File.join(@bucket, dirs, md5)
         filedir, contents, paths = bucket.class.paths(@bucket, md5)
 
         assert_equal(dir, filedir, "did not use a deeper file structure")