diff --git a/lib/puppet/functions/pty/spawn.rb b/lib/puppet/functions/pty/spawn.rb
index 991d9e0..1ea33a1 100644
--- a/lib/puppet/functions/pty/spawn.rb
+++ b/lib/puppet/functions/pty/spawn.rb
@@ -18,7 +18,9 @@
   #   The command to spawn.
   # @param block
   #   The code block, that is using PTY::IO object yielded to talk to the
-  #   command executed.
+  #   command executed. The process executed will be SIGTERM-ed and waited on
+  #   block exit.
+  # @return Block result.
   #
   # @example Spawn /bin/sh and get the hostname
   #   pty::spawn(['/bin/sh', '--norc']) |$pty| {
@@ -31,14 +33,14 @@
   dispatch :spawn_with_block do
     param 'Array[String[1]]', :cmd
     block_param 'Callable[PTY::IO]', :block
-    return_type 'Undef'
+    return_type 'Any'
   end
 
   # Spawns the specified command on a newly allocated pty (non-block form).
   #
   # @param cmd
   #   The command to spawn.
-  # @return The PTY::IO object
+  # @return The PTY::IO object.
   #
   # @example Spawn /bin/sh and get the hostname
   #   $pty = pty::spawn(['/bin/sh', '--norc'])
@@ -56,12 +58,10 @@
   def spawn_with_block(cmd)
     fail_if_not_in_plan
 
-    PTY.spawn(*cmd) do |input, output, pid|
-      pty_io = PuppetX::PTY::IO.new
-      pty_io.private_init(input, output, pid)
-      yield pty_io
-    end
-    nil
+    pty_io = spawn(cmd)
+    yield pty_io
+  ensure
+    pty_io&.close
   end
 
   def spawn(cmd)
diff --git a/lib/puppet_x/pty.rb b/lib/puppet_x/pty.rb
deleted file mode 100644
index e97ed36..0000000
--- a/lib/puppet_x/pty.rb
+++ /dev/null
@@ -1,4 +0,0 @@
-require 'puppet_x'
-
-# Define PuppetX::PTY namespace
-module PuppetX::PTY; end
diff --git a/lib/puppet_x/pty/io.rb b/lib/puppet_x/pty/io.rb
index 0296ada..f1a7073 100644
--- a/lib/puppet_x/pty/io.rb
+++ b/lib/puppet_x/pty/io.rb
@@ -1,10 +1,14 @@
 # frozen_string_literal: true
 
-require 'puppet_x/pty'
 require 'io/console'
 require 'io/wait'
 require 'expect'
 
+# Defined PuppetX namespace
+module PuppetX; end
+# Defined PuppetX::PTY namespace
+module PuppetX::PTY; end
+
 # Wrapper class for the pty interaction
 # @method write Write a string as is
 # @method puts Write a string with line feed appended
@@ -154,7 +158,7 @@ def type_in(msg, opts = {})
     msg.chars.each do |c|
       @output.putc c
       loop do
-        return nil unless @input.wait_readable(1) # Return Undef if timeout expired and no input here
+        return nil unless @input.wait_readable(30) # Return Undef if timeout expired and no input here
         i = @input.getc
         break if i == c
       end
diff --git a/spec/functions/pty/spawn_spec.rb b/spec/functions/pty/spawn_spec.rb
new file mode 100644
index 0000000..54a2d92
--- /dev/null
+++ b/spec/functions/pty/spawn_spec.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'bolt_spec/bolt_context'
+require 'pty'
+require_relative '../../../lib/puppet_x/pty/io'
+
+describe 'pty::spawn' do
+  let(:pty_io) { PuppetX::PTY::IO.new }
+  let(:command) { ['/nonexistent', '1', '2', '3'] }
+  let(:input) { instance_double(StringIO) }
+  let(:output) { instance_double(StringIO) }
+  let(:pid) { 123_456 }
+
+  include BoltSpec::BoltContext
+  # This function should always be tested in the Bolt context
+  around(:each) do |example|
+    in_bolt_context do
+      example.run
+    end
+  end
+
+  context 'in non-block form' do
+    it 'runs successfully' do
+      expect(PuppetX::PTY::IO).to receive(:new).and_return(pty_io)
+      expect(PTY).to receive(:spawn).with(*command).and_return([input, output, pid])
+      is_expected.to run.with_params(command).and_return(pty_io)
+    end
+  end
+end