2 files changed, 53 insertions, 1 deletions
diff --git a/spec/lib/request_spec.rb b/spec/lib/request_spec.rb
index 4d6b20dd5..939ac006a 100644
--- a/spec/lib/request_spec.rb
+++ b/spec/lib/request_spec.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'rails_helper'
+require 'securerandom'
 
 describe Request do
   subject { Request.new(:get, 'http://example.com') }
@@ -64,6 +65,12 @@ describe Request do
         expect_any_instance_of(HTTP::Client).to receive(:close)
         expect { |block| subject.perform &block }.to yield_control
       end
+
+      it 'returns response which implements body_with_limit' do
+        subject.perform do |response|
+          expect(response).to respond_to :body_with_limit
+        end
+      end
     end
 
     context 'with private host' do
@@ -81,4 +88,46 @@ describe Request do
       end
     end
   end
+
+  describe "response's body_with_limit method" do
+    it 'rejects body more than 1 megabyte by default' do
+      stub_request(:any, 'http://example.com').to_return(body: SecureRandom.random_bytes(2.megabytes))
+      expect { subject.perform { |response| response.body_with_limit } }.to raise_error Mastodon::LengthValidationError
+    end
+
+    it 'accepts body less than 1 megabyte by default' do
+      stub_request(:any, 'http://example.com').to_return(body: SecureRandom.random_bytes(2.kilobytes))
+      expect { subject.perform { |response| response.body_with_limit } }.not_to raise_error
+    end
+
+    it 'rejects body by given size' do
+      stub_request(:any, 'http://example.com').to_return(body: SecureRandom.random_bytes(2.kilobytes))
+      expect { subject.perform { |response| response.body_with_limit(1.kilobyte) } }.to raise_error Mastodon::LengthValidationError
+    end
+
+    it 'rejects too large chunked body' do
+      stub_request(:any, 'http://example.com').to_return(body: SecureRandom.random_bytes(2.megabytes), headers: { 'Transfer-Encoding' => 'chunked' })
+      expect { subject.perform { |response| response.body_with_limit } }.to raise_error Mastodon::LengthValidationError
+    end
+
+    it 'rejects too large monolithic body' do
+      stub_request(:any, 'http://example.com').to_return(body: SecureRandom.random_bytes(2.megabytes), headers: { 'Content-Length' => 2.megabytes })
+      expect { subject.perform { |response| response.body_with_limit } }.to raise_error Mastodon::LengthValidationError
+    end
+
+    it 'uses binary encoding if Content-Type does not tell encoding' do
+      stub_request(:any, 'http://example.com').to_return(body: '', headers: { 'Content-Type' => 'text/html' })
+      expect(subject.perform { |response| response.body_with_limit.encoding }).to eq Encoding::BINARY
+    end
+
+    it 'uses binary encoding if Content-Type tells unknown encoding' do
+      stub_request(:any, 'http://example.com').to_return(body: '', headers: { 'Content-Type' => 'text/html; charset=unknown' })
+      expect(subject.perform { |response| response.body_with_limit.encoding }).to eq Encoding::BINARY
+    end
+
+    it 'uses encoding specified by Content-Type' do
+      stub_request(:any, 'http://example.com').to_return(body: '', headers: { 'Content-Type' => 'text/html; charset=UTF-8' })
+      expect(subject.perform { |response| response.body_with_limit.encoding }).to eq Encoding::UTF_8
+    end
+  end
 end
diff --git a/spec/models/concerns/remotable_spec.rb b/spec/models/concerns/remotable_spec.rb
index 0b2dad23f..b39233739 100644
--- a/spec/models/concerns/remotable_spec.rb
+++ b/spec/models/concerns/remotable_spec.rb
@@ -29,7 +29,10 @@ RSpec.describe Remotable do
 
   context 'Remotable module is included' do
     before do
-      class Foo; include Remotable; end
+      class Foo
+        include Remotable
+        remotable_attachment :hoge, 1.kilobyte
+      end
     end
 
     let(:attribute_name) { "#{hoge}_remote_url".to_sym }