Bring in Shopify modifications

lib/tcr.rb

@@ -1,13 +1,16 @@
 require "tcr/cassette"
 require "tcr/configuration"
 require "tcr/errors"
-require "tcr/recordable_tcp_socket"
+require "tcr/recordable"
+require "tcr/socket_extension"
+require "tcr/ssl_socket_extension"
 require "tcr/version"
 require "socket"
 require "json"
 
-
 module TCR
+  ALL_PORTS = '*'
+
   extend self
 
   def configure
@@ -34,37 +37,31 @@ def disabled=(v)
     @disabled = v
   end
 
-  def save_session
+  def record_port?(port)
+    !disabled && configuration.hook_tcp_ports == ALL_PORTS || configuration.hook_tcp_ports.include?(port)
   end
 
   def use_cassette(name, options = {}, &block)
     raise ArgumentError, "`TCR.use_cassette` requires a block." unless block
-    TCR.cassette = Cassette.new(name)
-    yield
-    TCR.cassette = nil
+    begin
+      TCR.cassette = Cassette.build(name)
+      yield TCR.cassette
+    ensure
+      TCR.cassette.finish
+      TCR.cassette = nil
+    end
   end
 
   def turned_off(&block)
     raise ArgumentError, "`TCR.turned_off` requires a block." unless block
-    current_hook_tcp_ports = configuration.hook_tcp_ports
-    configuration.hook_tcp_ports = []
-    yield
-    configuration.hook_tcp_ports = current_hook_tcp_ports
-  end
-end
-
-
-# The monkey patch shim
-class TCPSocket
-  class << self
-    alias_method :real_open,  :open
-
-    def open(address, port)
-      if TCR.configuration.hook_tcp_ports.include?(port)
-        TCR::RecordableTCPSocket.new(address, port, TCR.cassette)
-      else
-        real_open(address, port)
-      end
+    begin
+      disabled = true
+      yield
+    ensure
+      disabled = false
     end
   end
 end
+
+TCPSocket.prepend(TCR::SocketExtension)
+OpenSSL::SSL::SSLSocket.send(:include, TCR::SSLSocketExtension)

lib/tcr/cassette.rb

@@ -2,39 +2,175 @@ module TCR
   class Cassette
     attr_reader :name
 
-    def initialize(name)
-      @name = name
-
-      if File.exists?(filename)
-        @recording = false
-        @contents = File.open(filename) { |f| f.read }
-        @sessions = JSON.parse(@contents)
+    def self.build(name)
+      if cassette_exists?(name)
+        RecordedCassette.new(name)
       else
-        @recording = true
-        @sessions = []
+        RecordingCassette.new(name)
       end
     end
 
-    def recording?
-      @recording
+    def self.filename(name)
+      "#{TCR.configuration.cassette_library_dir}/#{name}.json"
     end
 
-    def next_session
-      session = @sessions.shift
-      raise NoMoreSessionsError unless session
-      session
+    def self.cassette_exists?(name)
+      File.exists?(filename(name))
     end
 
-    def append(session)
-      raise "Can't append session unless recording" unless recording?
-      @sessions << session
-      File.open(filename, "w") { |f| f.write(JSON.pretty_generate(@sessions)) }
+    def initialize(name)
+      @name = name
     end
 
     protected
 
     def filename
-      "#{TCR.configuration.cassette_library_dir}/#{name}.json"
+      self.class.filename(name)
+    end
+
+    class RecordingCassette < Cassette
+      attr_reader :originally_recorded_at
+
+      def initialize(*args)
+        super
+        @originally_recorded_at = Time.now
+      end
+
+      def sessions
+        @sessions ||= []
+      end
+
+      def next_session
+        Session.new.tap do |session|
+          sessions << session
+        end
+      end
+
+      def finish
+        FileUtils.mkdir_p(File.dirname(filename))
+        File.open(filename, "w") { |f| f.write(serialized_form) }
+      end
+
+      private
+
+      def serialized_form
+        raw = {
+          "originally_recorded_at" => originally_recorded_at,
+          'sessions' => sessions.map(&:as_json)
+        }
+        JSON.pretty_generate(raw)
+      end
+
+      def empty?
+        true
+      end
+
+      class Session
+        def initialize
+          @recording = []
+        end
+
+        def connect(&block)
+          next_command('connect', ret: nil, &block)
+        end
+
+        def close(&block)
+          next_command('close', &block)
+        end
+
+        def read(*args, &block)
+          next_command('read', data: args, &block)
+        end
+
+        def write(str, &block)
+          next_command('write', data: str, &block)
+        end
+
+        def as_json
+          {"recording" => @recording}
+        end
+
+        private
+
+        def next_command(command, options={}, &block)
+          yield.tap do |return_value|
+            return_value = options[:ret] if options.has_key?(:ret)
+            @recording << [command, return_value] + Array(options[:data])
+          end
+        end
+      end
+    end
+
+    class RecordedCassette < Cassette
+      def sessions
+        @sessions ||= serialized_form['sessions'].map{|raw| Session.new(raw)}
+      end
+
+      def originally_recorded_at
+        serialized_form['originally_recorded_at']
+      end
+
+      def next_session
+        session = sessions.shift
+        raise NoMoreSessionsError unless session
+        session
+      end
+
+      def finish
+        # no-op
+      end
+
+      def empty?
+        sessions.all?(&:empty?)
+      end
+
+      private
+
+      def serialized_form
+        @serialized_form ||= begin
+          raw = File.open(filename) { |f| f.read }
+          JSON.parse(raw)
+        end
+      end
+
+      class Session
+        def initialize(raw)
+          @recording = raw["recording"]
+        end
+
+        def connect
+          next_command('connect')
+        end
+
+        def close
+          next_command('close')
+        end
+
+        def read(*args)
+          next_command('read') do |str, *data|
+            raise TCR::DataMismatchError.new("Expected to read to be called with args '#{args.inspect}' but was called with '#{data.inspect}'") unless args == data
+          end
+        end
+
+        def write(str)
+          next_command('write') do |len, data|
+            raise TCR::DataMismatchError.new("Expected to write '#{str}' but next data in recording was '#{data}'") unless str == data
+          end
+        end
+
+        def empty?
+          @recording.empty?
+        end
+
+        private
+
+        def next_command(expected)
+          command, return_value, data = @recording.shift
+          raise TCR::CommandMismatchError.new("Expected to '#{expected}' but next in recording was '#{command}'") unless expected == command
+          yield return_value, *data if block_given?
+          return_value
+        end
+      end
     end
   end
-end
+end

lib/tcr/errors.rb

@@ -1,6 +1,6 @@
 module TCR
   class TCRError < StandardError; end
-  class NoCassetteError < TCRError; end
   class NoMoreSessionsError < TCRError; end
-  class DirectionMismatchError < TCRError; end
+  class CommandMismatchError < TCRError; end
+  class DataMismatchError < TCRError; end
 end

lib/tcr/recordable.rb

@@ -0,0 +1,45 @@
+module TCR
+  module Recordable
+    attr_accessor :cassette
+
+    def recording
+      @recording ||= cassette.next_session
+    end
+
+    def connect
+      recording.connect do
+        super
+      end
+    end
+
+    def read_nonblock(bytes)
+      recording.read do
+        super
+      end
+    end
+
+    def gets(*args)
+      recording.read do
+        super
+      end
+    end
+
+    def write(str)
+      recording.write(str) do
+        super
+      end
+    end
+
+    def read(*args)
+      recording.read do
+        super
+      end
+    end
+
+    def close
+      recording.close do
+        super
+      end
+    end
+  end
+end