lib/github/kv.rb in github-ds-0.2.10 vs lib/github/kv.rb in github-ds-0.2.11

@@ -49,12 +49,28 @@
     ValueLengthError = Class.new(StandardError)
     UnavailableError = Class.new(StandardError)
 
     class MissingConnectionError < StandardError; end
 
-    def initialize(encapsulated_errors = [SystemCallError], &conn_block)
+    attr_accessor :use_local_time
+
+    # initialize :: [Exception], Boolean, Proc -> nil
+    #
+    # Initialize a new KV instance.
+    #
+    # encapsulated_errors - An Array of Exception subclasses that, when raised,
+    #                       will be replaced with UnavailableError.
+    # use_local_time:     - Whether to use Ruby's `Time.now` instaed of MySQL's
+    #                       `NOW()` function. This is mostly useful in testing
+    #                       where time needs to be modified (eg. Timecop).
+    #                       Default false.
+    # &conn_block         - A block to call to open a new database connection.
+    #
+    # Returns nothing.
+    def initialize(encapsulated_errors = [SystemCallError], use_local_time: false, &conn_block)
       @encapsulated_errors = encapsulated_errors
+      @use_local_time = use_local_time
       @conn_block = conn_block
     end
 
     def connection
       @conn_block.try(:call) || (raise MissingConnectionError, "KV must be initialized with a block that returns a connection")
@@ -91,12 +107,12 @@
     #
     def mget(keys)
       validate_key_array(keys)
 
       Result.new {
-        kvs = GitHub::SQL.results(<<-SQL, :keys => keys, :connection => connection).to_h
-          SELECT `key`, value FROM key_values WHERE `key` IN :keys AND (`expires_at` IS NULL OR `expires_at` > NOW())
+        kvs = GitHub::SQL.results(<<-SQL, :keys => keys, :now => now, :connection => connection).to_h
+          SELECT `key`, value FROM key_values WHERE `key` IN :keys AND (`expires_at` IS NULL OR `expires_at` > :now)
         SQL
 
         keys.map { |key| kvs[key] }
       }
     end
@@ -135,11 +151,11 @@
       validate_key_value_hash(kvs)
       validate_expires(expires) if expires
 
       rows = kvs.map { |key, value|
         value = value.is_a?(GitHub::SQL::Literal) ? value : GitHub::SQL::BINARY(value)
-        [key, value, GitHub::SQL::NOW, GitHub::SQL::NOW, expires || GitHub::SQL::NULL]
+        [key, value, now, now, expires || GitHub::SQL::NULL]
       }
 
       encapsulate_error do
         GitHub::SQL.run(<<-SQL, :rows => GitHub::SQL::ROWS(rows), :connection => connection)
           INSERT INTO key_values (`key`, value, created_at, updated_at, expires_at)
@@ -184,12 +200,12 @@
     #
     def mexists(keys)
       validate_key_array(keys)
 
       Result.new {
-        existing_keys = GitHub::SQL.values(<<-SQL, :keys => keys, :connection => connection).to_set
-          SELECT `key` FROM key_values WHERE `key` IN :keys AND (`expires_at` IS NULL OR `expires_at` > NOW())
+        existing_keys = GitHub::SQL.values(<<-SQL, :keys => keys, :now => now, :connection => connection).to_set
+          SELECT `key` FROM key_values WHERE `key` IN :keys AND (`expires_at` IS NULL OR `expires_at` > :now)
         SQL
 
         keys.map { |key| existing_keys.include?(key) }
       }
     end
@@ -220,18 +236,18 @@
       encapsulate_error {
         # if the key already exists but has expired, prune it first. We could
         # achieve the same thing with the right INSERT ... ON DUPLICATE KEY UPDATE
         # query, but then we would not be able to rely on affected_rows
 
-        GitHub::SQL.run(<<-SQL, :key => key, :connection => connection)
-          DELETE FROM key_values WHERE `key` = :key AND expires_at <= NOW()
+        GitHub::SQL.run(<<-SQL, :key => key, :now => now, :connection => connection)
+          DELETE FROM key_values WHERE `key` = :key AND expires_at <= :now
         SQL
 
         value = value.is_a?(GitHub::SQL::Literal) ? value : GitHub::SQL::BINARY(value)
-        sql = GitHub::SQL.run(<<-SQL, :key => key, :value => value, :expires => expires || GitHub::SQL::NULL, :connection => connection)
+        sql = GitHub::SQL.run(<<-SQL, :key => key, :value => value, :now => now, :expires => expires || GitHub::SQL::NULL, :connection => connection)
           INSERT IGNORE INTO key_values (`key`, value, created_at, updated_at, expires_at)
-          VALUES (:key, :value, NOW(), NOW(), :expires)
+          VALUES (:key, :value, :now, :now, :expires)
         SQL
 
         sql.affected_rows > 0
       }
     end
@@ -286,17 +302,21 @@
     #
     def ttl(key)
       validate_key(key)
 
       Result.new {
-        GitHub::SQL.value(<<-SQL, :key => key, :connection => connection)
+        GitHub::SQL.value(<<-SQL, :key => key, :now => now, :connection => connection)
           SELECT expires_at FROM key_values
-          WHERE `key` = :key AND (expires_at IS NULL OR expires_at > NOW())
+          WHERE `key` = :key AND (expires_at IS NULL OR expires_at > :now)
         SQL
       }
     end
 
   private
+    def now
+      use_local_time ? Time.now : GitHub::SQL::NOW
+    end
+
     def validate_key(key, error_message: nil)
       unless key.is_a?(String)
         raise TypeError, error_message || "key must be a String in #{self.class.name}, but was #{key.class}"
       end