Refactor connect/disconnect logic

bmhughes · Jan 19, 2024 · fe0b804 · fe0b804
1 parent ac5e272
commit fe0b804
Showing 1 changed file with 17 additions and 10 deletions.
diff --git a/lib/tp_link_smartplug/device.rb b/lib/tp_link_smartplug/device.rb
@@ -15,9 +15,9 @@ class Device < TpLinkSmartplug::Base
     include TpLinkSmartplug::Helpers
     include TpLinkSmartplug::Message
 
-    attr_accessor :address, :port, :timeout, :poll_auto_close, :debug
+    attr_accessor :address, :port, :timeout, :poll_auto_close, :auto_connect, :debug
 
-    def initialize(address:, port: 9999)
+    def initialize(address:, port: 9999, auto_connect: true, auto_disconnect: true)
       super()
 
       @address = IPAddr.new(address, Socket::AF_INET)
@@ -26,7 +26,10 @@ def initialize(address:, port: 9999)
       @debug = false
       @sockaddr = Addrinfo.getaddrinfo(@address.to_s, @port, Socket::PF_INET, :STREAM, 6).first.to_sockaddr
       @socket = Socket.new(Socket::AF_INET, Socket::SOCK_STREAM)
-      @poll_auto_close = true
+      @poll_auto_close = auto_disconnect
+      @auto_connect = auto_connect
+
+      connect if auto_connect
     end
 
     # Open connection to plug
@@ -41,7 +44,7 @@ def connect
         @socket.connect_nonblock(@sockaddr)
         debug_message('Connected') if @debug
       rescue IO::WaitWritable
-        if @socket.wait_writable(timeout)
+        if @socket.wait_writable(@timeout)
           begin
             @socket.connect_nonblock(@sockaddr)
           rescue Errno::EISCONN
@@ -55,7 +58,7 @@ def connect
             raise
           end
         else
-          @socket.close
+          disconnect
           raise TpLinkSmartplug::DeviceError, "Connection timeout connecting to address #{@address}, port #{@port}."
         end
       rescue Errno::EISCONN
@@ -71,7 +74,7 @@ def connect
     # Close connection to plug
     #
     def disconnect
-      @socket.close unless @socket.closed?
+      @socket.close unless closed?
     end
 
     alias_method :close, :disconnect
@@ -87,15 +90,19 @@ def open?
     #
     # @return [True, False]
     def closed?
-      @socket.closed?
+      @socket.recvfrom_nonblock(0)
+    rescue IO::WaitReadable
+      false
+    rescue Errno::ENOTCONN, IOError
+      true
     end
 
     # Polls plug with a command
     #
     # @param command [String] the command to send to the plug
     # @return [Hash] the output from the plug command
     def poll(command:)
-      connect
+      connect if closed?
 
       begin
         debug_message("Sending: #{decrypt(encrypt(command)[4..(command.length + 4)])}") if @debug
@@ -114,10 +121,10 @@ def poll(command:)
 
       if @poll_auto_close && !closed?
         disconnect
-        raise 'Error occured during disconnect' unless closed?
+        raise DeviceError, 'Error occured during disconnect' unless closed?
       end
 
-      raise 'No data received' if nil_or_empty?(data)
+      raise DeviceError, 'No data received' if nil_or_empty?(data)
 
       debug_message("Received Raw: #{data.split('\\')}") if @debug
       data = decrypt(data[4..data.length])