gistfile1.rb

diff --git a/lib/mock_server.rb b/lib/mock_server.rb
index d09a2b4..99c4089 100644
--- a/lib/mock_server.rb
+++ b/lib/mock_server.rb
@@ -6,11 +6,9 @@ class MockServer
     use Rack::ShowExceptions
   end
 
-  def initialize(port = 4000, &block)
+  def initialize(app, port = 4000, &block)
+    @app = app
     @port = port
-
-    @app = Class.new(Sinatra::Base)
-    @app.class_eval(&block)
   end
 
   def start
@@ -27,7 +25,9 @@ class MockServer
 
   module Methods
     def mock_server(*args, &block)
-      @server = MockServer.new(*args, &block).start
+      app = Class.new(Sinatra::Base)
+      app.class_eval(&block)
+      @server = MockServer.new(app, *args, &block).start
     end
   end
 
diff --git a/test/mock_server_test.rb b/test/mock_server_test.rb
index f1cfbf5..b6097f5 100644
--- a/test/mock_server_test.rb
+++ b/test/mock_server_test.rb
@@ -4,15 +4,16 @@ require "test/unit"
 require "ruby-debug"
 require "open-uri"
 
-class MockServerTest < Test::Unit::TestCase
-  @@server = MockServer.new do
-    get "/" do
-      "Hello"
-    end
+class HelloWorldSinatra < Sinatra::Base
+  get "/" do
+    "Hello"
   end
+end
 
+class MockServerTest < Test::Unit::TestCase
   def setup
-    @@server.start
+    @server = MockServer.new(HelloWorldSinatra)
+    @server.start
   end
 
   def test_server
@@ -20,16 +21,32 @@ class MockServerTest < Test::Unit::TestCase
   end
 end
 
+HelloWorldRackBuilder = Rack::Builder.new do
+  run lambda {|env|
+    [200, {"Content-Type" => "text/plain", "Content-Length" => "7"}, ["Rackup!"]]
+  }
+end
+
+class MockServerRackBuilderTest < Test::Unit::TestCase
+  def setup
+    @server = MockServer.new(HelloWorldRackBuilder, 4001)
+    @server.start
+  end
+
+  def test_server
+    assert_equal "Rackup!", open("http://localhost:4001").read
+  end end
+
 class MockServerMethodsTest < Test::Unit::TestCase
   extend MockServer::Methods
 
-  mock_server(4001) {
+  mock_server(4002) {
     get "/" do
       "Goodbye"
     end
   }
 
   def test_server
-    assert_equal "Goodbye", open("http://localhost:4001").read
+    assert_equal "Goodbye", open("http://localhost:4002").read
   end
 end