Add Reline.test_mode to use with special I/O

2 files changed, 31 insertions, 11 deletions

diff --git a/lib/reline.rb b/lib/reline.rb
index 61b94d727a..a5acc52f03 100644
--- a/lib/reline.rb
+++ b/lib/reline.rb
@@ -180,12 +180,19 @@ module Reline
     @@line_editor&.delete_text(start, length)
   end
 
+  private_class_method def self.test_mode
+    remove_const('IOGate') if const_defined?('IOGate')
+    const_set('IOGate', Reline::GeneralIO)
+  end
+
   def self.input=(val)
     raise TypeError unless val.respond_to?(:getc) or val.nil?
     if val.respond_to?(:getc)
-      Reline::GeneralIO.input = val
-      remove_const('IOGate') if const_defined?('IOGate')
-      const_set('IOGate', Reline::GeneralIO)
+      if defined?(Reline::ANSI) and IOGate == Reline::ANSI
+        Reline::ANSI.input = val
+      elsif IOGate == Reline::GeneralIO
+        Reline::GeneralIO.input = val
+      end
     end
   end
 
@@ -193,6 +200,9 @@ module Reline
   def self.output=(val)
     raise TypeError unless val.respond_to?(:write) or val.nil?
     @@output = val
+    if defined?(Reline::ANSI) and IOGate == Reline::ANSI
+      Reline::ANSI.output = val
+    end
   end
 
   def self.vi_editing_mode
diff --git a/lib/reline/ansi.rb b/lib/reline/ansi.rb
index c0637139f7..34720d352c 100644
--- a/lib/reline/ansi.rb
+++ b/lib/reline/ansi.rb
@@ -1,23 +1,33 @@
 class Reline::ANSI
+  @@input = STDIN
+  def self.input=(val)
+    @@input = val
+  end
+
+  @@output = STDOUT
+  def self.output=(val)
+    @@output = val
+  end
+
   def self.getc
     c = nil
     loop do
-      result = select([STDIN], [], [], 0.1)
+      result = select([@@input], [], [], 0.1)
       next if result.nil?
-      c = STDIN.read(1)
+      c = @@input.read(1)
       break
     end
     c&.ord
   end
 
   def self.get_screen_size
-    STDIN.winsize
+    @@input.winsize
   rescue Errno::ENOTTY
     [24, 80]
   end
 
   def self.set_screen_size(rows, columns)
-    STDIN.winsize = [rows, columns]
+    @@input.winsize = [rows, columns]
     self
   rescue Errno::ENOTTY
     self
@@ -26,9 +36,9 @@ class Reline::ANSI
   def self.cursor_pos
     begin
       res = ''
-      STDIN.raw do |stdin|
-        STDOUT << "\e[6n"
-        STDOUT.flush
+      @@input.raw do |stdin|
+        @@output << "\e[6n"
+        @@output.flush
         while (c = stdin.getc) != 'R'
           res << c if c
         end
@@ -37,7 +47,7 @@ class Reline::ANSI
       column = m[:column].to_i - 1
       row = m[:row].to_i - 1
     rescue Errno::ENOTTY
-      buf = STDOUT.pread(STDOUT.pos, 0)
+      buf = @@output.pread(@@output.pos, 0)
       row = buf.count("\n")
       column = buf.rindex("\n") ? (buf.size - buf.rindex("\n")) - 1 : 0
     end