6 files changed, 20 insertions, 38 deletions
diff --git a/lib/chotto.rb b/lib/chotto.rb
index 6e1f908..71a1ac8 100644
--- a/lib/chotto.rb
+++ b/lib/chotto.rb
@@ -9,7 +9,6 @@ require_relative 'chotto/helpers'
 require_relative 'chotto/message'
 require_relative 'chotto/messages'
 require_relative 'chotto/ruleset'
-require_relative 'chotto/tags'
 
 module Chotto
   class << self
@@ -18,15 +17,16 @@ module Chotto
 
     def configure(&block)
       @config ||= Config.new
+      @config.db_class = ::Notmuch::Database
       @rule_sets = []
 
       instance_eval(&block)
 
-      @db = Database.new(path: config.database_path)
+      @db = Database.new(path: config.database_path, db_class: config.db_class)
     end
 
     def rule_set(name, &block)
-      @rule_sets << RuleSet.new(name, block)
+      @rule_sets << RuleSet.new(name, db, block)
     end
 
     def close_db
@@ -36,8 +36,6 @@ module Chotto
 end
 
 def eval_rules
-  Chotto.rule_sets.each do |rule_set|
-    rule_set.run
-  end
+  Chotto.rule_sets.each(&:run)
   Chotto.close_db
 end
diff --git a/lib/chotto/config.rb b/lib/chotto/config.rb
index 267e72c..2edfcc8 100644
--- a/lib/chotto/config.rb
+++ b/lib/chotto/config.rb
@@ -2,6 +2,6 @@
 
 module Chotto
   class Config
-    attr_accessor :database_path
+    attr_accessor :database_path, :db_class
   end
 end
diff --git a/lib/chotto/database.rb b/lib/chotto/database.rb
index 893c970..a197588 100644
--- a/lib/chotto/database.rb
+++ b/lib/chotto/database.rb
@@ -4,16 +4,12 @@ module Chotto
   class Database
     attr_reader :db
 
-    def initialize(path:)
-      @db = ::Notmuch::Database.new(path, mode: Notmuch::MODE_READ_WRITE)
-    end
-
-    def query(query)
-      db.query(query)
+    def initialize(path:, db_class:)
+      @db = db_class.new(path, mode: Notmuch::MODE_READ_WRITE)
     end
 
     def search_messages(query)
-      query(query).search_messages
+      db.query(query).search_messages
     end
 
     def close
diff --git a/lib/chotto/message.rb b/lib/chotto/message.rb
index ff382f8..282a030 100644
--- a/lib/chotto/message.rb
+++ b/lib/chotto/message.rb
@@ -11,18 +11,20 @@ module Chotto
     end
 
     def method_missing(method_name, *_args)
-      handle_get(Chotto::Helpers.header_name_from_dsl(method_name))
-    end
-
-    def handle_get(header_name)
-      message.header(header_name) if message.header(header_name)
+      handle_get_header(Chotto::Helpers.header_name_from_dsl(method_name))
     end
 
     def save!
       message.remove_all_tags
       tags.each do |tag|
-        message.add_tag(tag)
+        message.add_tag(tag.to_s)
       end
     end
+
+    private
+
+    def handle_get_header(header_name)
+      message.header(header_name)
+    end
   end
 end
diff --git a/lib/chotto/ruleset.rb b/lib/chotto/ruleset.rb
index 289aa2f..08c9dbd 100644
--- a/lib/chotto/ruleset.rb
+++ b/lib/chotto/ruleset.rb
@@ -1,7 +1,9 @@
+# frozen_string_literal: true
+
 module Chotto
-  RuleSet = Struct.new(:name, :rule) do
+  RuleSet = Struct.new(:name, :db, :rule) do
     def messages
-      Messages.new(db: Chotto.db)
+      Messages.new(db: db)
     end
 
     def run
diff --git a/lib/chotto/tags.rb b/lib/chotto/tags.rb
deleted file mode 100644
index 52caeec..0000000
--- a/lib/chotto/tags.rb
+++ /dev/null
@@ -1,16 +0,0 @@
-module Chotto
-  class Tags
-    include Enumerable
-
-    attr_accessor :tags, :message
-
-    def initialize(tags:, message:)
-      @tags = tags
-      @messae = message
-    end
-
-    def each(&block)
-      tags.each(&block)
-    end
-  end
-end