Merge pull request #546 from neektza/injectable_connection_sockets

Enable injection of custom TCP/SSL socket classes.

lib/twitter/streaming/client.rb

@@ -9,13 +9,17 @@ module Twitter
   module Streaming
     class Client < Twitter::Client
       attr_writer :connection
+      attr_accessor :tcp_socket_class, :ssl_socket_class
 
       # Initializes a new Client object
       #
+      # @param options [Hash] A customizable set of options.
+      # @option options [String] :tcp_socket_class A class that Connection will use to create a new TCP socket.
+      # @option options [String] :ssl_socket_class A class that Connection will use to create a new SSL socket.
       # @return [Twitter::Streaming::Client]
       def initialize(options = {})
         super
-        @connection = Streaming::Connection.new
+        @connection = Streaming::Connection.new(options)
       end
 
       # Returns public statuses that match one or more filter predicates

lib/twitter/streaming/connection.rb

@@ -5,10 +5,17 @@
 module Twitter
   module Streaming
     class Connection
-      def stream(request, response)
+      def initialize(opts = {})
+        @tcp_socket_class = opts.fetch(:tcp_socket_class) { TCPSocket }
+        @ssl_socket_class = opts.fetch(:ssl_socket_class) { OpenSSL::SSL::SSLSocket }
+      end
+      attr_reader :tcp_socket_class, :ssl_socket_class
+
+      def stream(request, response, opts = {})
         client_context = OpenSSL::SSL::SSLContext.new
-        client         = TCPSocket.new(Resolv.getaddress(request.uri.host), request.uri.port)
-        ssl_client     = OpenSSL::SSL::SSLSocket.new(client, client_context)
+        client         = @tcp_socket_class.new(Resolv.getaddress(request.uri.host), request.uri.port)
+        ssl_client     = @ssl_socket_class.new(client, client_context)
+
         ssl_client.connect
         request.stream(ssl_client)
         while body = ssl_client.readpartial(1024) # rubocop:disable AssignmentInCondition, WhileUntilModifier

spec/twitter/streaming/connection_spec.rb

@@ -0,0 +1,32 @@
+require 'helper'
+
+describe Twitter::Streaming::Connection do
+
+  describe 'initialize' do
+    context 'no options provided' do
+      subject(:connection) { Twitter::Streaming::Connection.new }
+
+      it 'sets the default socket classes' do
+        expect(connection.tcp_socket_class).to eq TCPSocket
+        expect(connection.ssl_socket_class).to eq OpenSSL::SSL::SSLSocket
+      end
+    end
+
+    context 'custom socket classes provided in opts' do
+      class DummyTCPSocket; end
+      class DummySSLSocket; end
+
+      subject(:connection) do
+        Twitter::Streaming::Connection.new(
+          :tcp_socket_class => DummyTCPSocket,
+          :ssl_socket_class => DummySSLSocket
+        )
+      end
+
+      it 'sets the default socket classes' do
+        expect(connection.tcp_socket_class).to eq DummyTCPSocket
+        expect(connection.ssl_socket_class).to eq DummySSLSocket
+      end
+    end
+  end
+end