diff --git a/app/models/media_asset.rb b/app/models/media_asset.rb index a683d061b..890b6b333 100644 --- a/app/models/media_asset.rb +++ b/app/models/media_asset.rb @@ -43,7 +43,7 @@ class MediaAsset < ApplicationRecord failed: 500, } - validates :md5, uniqueness: { conditions: -> { where(status: [:processing, :active]) } } + validates :md5, uniqueness: { conditions: -> { where(status: [:processing, :active]) } }, if: :md5_changed? validates :file_ext, inclusion: { in: FILE_TYPES, message: "File is not an image or video" } validates :file_size, numericality: { less_than_or_equal_to: Danbooru.config.max_file_size, message: ->(asset, _) { "too large (size: #{asset.file_size.to_formatted_s(:human_size)}; max size: #{Danbooru.config.max_file_size.to_formatted_s(:human_size)})" } } validates :file_key, length: { is: FILE_KEY_LENGTH }, uniqueness: true, if: :file_key_changed? @@ -232,6 +232,13 @@ class MediaAsset < ApplicationRecord media_asset = create!(file: media_file, status: :processing) yield media_asset if block_given? + + # XXX should do this in parallel with thumbnail generation. + # XXX shouldn't generate thumbnail twice (very slow for ugoira) + media_asset.update!(ai_tags: media_file.preview(360, 360).ai_tags) + media_asset.update!(pixiv_ugoira_frame_data: PixivUgoiraFrameData.new(data: media_file.frame_data, content_type: "image/jpeg")) if media_asset.is_ugoira? + media_asset.update!(media_metadata: MediaMetadata.new(file: media_file)) + media_asset.distribute_files!(media_file) media_asset.update!(status: :active) media_asset @@ -278,9 +285,6 @@ class MediaAsset < ApplicationRecord self.image_width = media_file.width self.image_height = media_file.height self.duration = media_file.duration - self.media_metadata = MediaMetadata.new(file: media_file) - self.pixiv_ugoira_frame_data = PixivUgoiraFrameData.new(data: media_file.frame_data, content_type: "image/jpeg") if is_ugoira? - self.ai_tags = media_file.preview(360, 360).ai_tags # XXX should do this in parallel with thumbnail generation. end def regenerate_ai_tags! diff --git a/test/functional/uploads_controller_test.rb b/test/functional/uploads_controller_test.rb index 0524b2b29..d9cbc720f 100644 --- a/test/functional/uploads_controller_test.rb +++ b/test/functional/uploads_controller_test.rb @@ -245,6 +245,19 @@ class UploadsControllerTest < ActionDispatch::IntegrationTest assert_equal(source2, upload.source) end + should "save the AI tags" do + mock_autotagger_evaluate({ "1girl": 0.5 }) + upload = assert_successful_upload("test/files/test.jpg") + + assert_equal(1, upload.media_assets.first.ai_tags.count) + end + + should "save the EXIF metadata" do + upload = assert_successful_upload("test/files/test.jpg") + + assert_equal(true, upload.media_assets.first.media_metadata.present?) + end + context "uploading a file from your computer" do should_upload_successfully("test/files/test.jpg") should_upload_successfully("test/files/test.png") diff --git a/test/test_helper.rb b/test/test_helper.rb index d96e278fe..d5d74871a 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -23,6 +23,7 @@ class ActiveSupport::TestCase extend PostArchiveTestHelper extend PoolArchiveTestHelper include ReportbooruHelper + include AutotaggerHelper include DownloadTestHelper include IqdbTestHelper include UploadTestHelper diff --git a/test/test_helpers/autotagger_helper.rb b/test/test_helpers/autotagger_helper.rb new file mode 100644 index 000000000..b08a605b9 --- /dev/null +++ b/test/test_helpers/autotagger_helper.rb @@ -0,0 +1,11 @@ +module AutotaggerHelper + def mock_autotagger_evaluate(tags, http: Danbooru::Http.any_instance) + tags.keys.each { |name| create(:tag, name: name) } + + Danbooru.config.stubs(:autotagger_url).returns("http://localhost:5000") + body = [{ filename: "test.jpg", tags: tags }] + + response = HTTP::Response.new(status: 200, body: body.to_json, version: "1.1", request: nil, headers: { "Content-Type": "application/json" }) + http.stubs(:post).with("http://localhost:5000/evaluate", anything).returns(response) + end +end