From 3a6865e172914c97ecf2651bf11302f5e1beb1cc Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Fri, 8 Jan 2010 19:11:49 -0500
Subject: [PATCH] minor: db class cleanup

---
 lib/mongo/db.rb | 31 ++++++++++++++++++++-----------
 test/test_db.rb |  4 ++--
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 2f2484b..03549cf 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -21,7 +21,7 @@ require 'thread'
 
 module Mongo
 
-  # A Mongo database.
+  # A MongoDB database.
   class DB
 
     SYSTEM_NAMESPACE_COLLECTION = "system.namespaces"
@@ -55,15 +55,7 @@ module Mongo
     # The logger instance if :logger is passed to initialize.
     attr_reader :logger
 
-    # The primary key factory object (or +nil+).
-    attr_reader :pk_factory
-
-    def pk_factory=(pk_factory)
-      raise "error: can not change PK factory" if @pk_factory
-      @pk_factory = pk_factory
-    end
-
-    # Instances of DB are normally obtained by calling Mongo#db.
+        # Instances of DB are normally obtained by calling Mongo#db.
     #
     # db_name :: The database name
     #
@@ -122,7 +114,8 @@ module Mongo
     # Deauthorizes use for this database for this connection.
     def logout
       doc = command(:logout => 1)
-      raise "error logging out: #{doc.inspect}" unless ok?(doc)
+      return true if ok?(doc)
+      raise "error logging out: #{doc.inspect}"
     end
 
     # Returns an array of collection names in this database.
@@ -378,6 +371,22 @@ module Mongo
       "#{@name}.#{collection_name}"
     end
 
+    # The primary key factory object (or +nil+).
+    def pk_factory
+      @pk_factory
+    end
+
+    # Specify a primary key factory if not already set.
+    #
+    # @raise [MongoArgumentError] if the primary key factory has already been set.
+    def pk_factory=(pk_factory)
+      if @pk_factory
+        raise MongoArgumentError, "Cannot change primary key factory once it's been set"
+      end
+
+      @pk_factory = pk_factory
+    end
+
     # Return the current database profiling level.
     #
     # @return [Symbol] :off, :slow_only, or :all
diff --git a/test/test_db.rb b/test/test_db.rb
index 1ea79ba..f8c6fd8 100644
--- a/test/test_db.rb
+++ b/test/test_db.rb
@@ -132,7 +132,7 @@ class DBTest < Test::Unit::TestCase
       db.pk_factory = Object.new
       fail "error: expected exception"
     rescue => ex
-      assert_match /can not change PK factory/, ex.to_s
+      assert_match /Cannot change/, ex.to_s
     ensure
       conn.close
     end
@@ -145,7 +145,7 @@ class DBTest < Test::Unit::TestCase
   end
 
   def test_logout
-    @@db.logout                  # only testing that we don't throw exception
+    assert @@db.logout
   end
 
   def test_error