@@ -124,6 +124,9 @@ const MediaReplaceFlow = ( {
 			onFileChange: ( [ media ] ) => {
 				selectMedia( media, closeMenu );
 			},
+			onSuccess: ( [ media ] ) => {
+				selectMedia( media, closeMenu );
+			},
 			onError: onUploadError,
 		} );
 	};
 
@@ -1,15 +1,17 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { useState } from '@wordpress/element';
 import { SlotFillProvider } from '@wordpress/components';
+import { dispatch } from '@wordpress/data';
 import MediaReplaceFlow from '../';
+import { store as blockEditorStore } from '../../../store';
 
 globalThis.wpVitest.mockMatchMedia();
 
 const noop = () => {};
 
-function TestWrapper() {
+function TestWrapper( { onSelect = noop, ...props } ) {
 	const [ mediaURL, setMediaURL ] = useState( 'https://example.media' );
 	return (
 		<SlotFillProvider>
@@ -18,10 +20,11 @@ function TestWrapper() {
 				mediaURL={ mediaURL }
 				allowedTypes={ [ 'png' ] }
 				accept="image/*"
-				onSelect={ noop }
+				onSelect={ onSelect }
 				onSelectURL={ setMediaURL }
 				onError={ noop }
 				onCloseModal={ noop }
+				{ ...props }
 			/>
 		</SlotFillProvider>
 	);
@@ -124,4 +127,40 @@ describe( 'General media replace flow', () => {
 			} )
 		).toHaveAttribute( 'href', 'https://new.example.media' );
 	} );
+
+	it( 'calls onSelect when mediaUpload completes via onSuccess', async () => {
+		const user = userEvent.setup();
+		const onSelect = vi.fn();
+		const testMedia = {
+			id: 42,
+			url: 'https://example.com/updated-file.png',
+			title: 'updated-file',
+		};
+
+		dispatch( blockEditorStore ).updateSettings( {
+			mediaUpload: vi.fn( ( { onFileChange, onSuccess } ) => {
+				onFileChange( [ { url: 'blob:https://example.com/uuid' } ] );
+				onSuccess( [ testMedia ] );
+			} ),
+		} );
+
+		render( <TestWrapper onSelect={ onSelect } /> );
+
+		await user.click(
+			screen.getByRole( 'button', {
+				expanded: false,
+				name: 'Replace',
+			} )
+		);
+
+		const uploadInput = screen.getByTestId( 'form-file-upload-input' );
+		const file = new File( [ 'sample' ], 'test.png', {
+			type: 'image/png',
+		} );
+		await user.upload( uploadInput, file );
+
+		await waitFor( () => {
+			expect( onSelect ).toHaveBeenLastCalledWith( testMedia );
+		} );
+	} );
 } );