Add a test and make it actually work

SJSAssetExportSession.xcodeproj/project.pbxproj

@@ -13,6 +13,7 @@
 		7B9BC0192C305D2C00C160C2 /* SJSAssetExportSessionTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */; };
 		7B9BC01A2C305D2C00C160C2 /* SJSAssetExportSession.h in Headers */ = {isa = PBXBuildFile; fileRef = 7B9BC00C2C305D2C00C160C2 /* SJSAssetExportSession.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		7B9BC0282C30612C00C160C2 /* ExportSession.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B9BC0272C30612C00C160C2 /* ExportSession.swift */; };
+		7BC5FC712C3A52B50090B757 /* test.mov in Resources */ = {isa = PBXBuildFile; fileRef = 7BC5FC702C3A52B50090B757 /* test.mov */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXContainerItemProxy section */
@@ -33,6 +34,7 @@
 		7B9BC0132C305D2C00C160C2 /* SJSAssetExportSessionTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = SJSAssetExportSessionTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SJSAssetExportSessionTests.swift; sourceTree = "<group>"; };
 		7B9BC0272C30612C00C160C2 /* ExportSession.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ExportSession.swift; sourceTree = "<group>"; };
+		7BC5FC702C3A52B50090B757 /* test.mov */ = {isa = PBXFileReference; lastKnownFileType = video.quicktime; path = test.mov; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -86,11 +88,20 @@
 		7B9BC0172C305D2C00C160C2 /* SJSAssetExportSessionTests */ = {
 			isa = PBXGroup;
 			children = (
+				7BC5FC6D2C3A525A0090B757 /* Resources */,
 				7B9BC0182C305D2C00C160C2 /* SJSAssetExportSessionTests.swift */,
 			);
 			path = SJSAssetExportSessionTests;
 			sourceTree = "<group>";
 		};
+		7BC5FC6D2C3A525A0090B757 /* Resources */ = {
+			isa = PBXGroup;
+			children = (
+				7BC5FC702C3A52B50090B757 /* test.mov */,
+			);
+			path = Resources;
+			sourceTree = "<group>";
+		};
 /* End PBXGroup section */
 
 /* Begin PBXHeadersBuildPhase section */
@@ -190,6 +201,7 @@
 			isa = PBXResourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				7BC5FC712C3A52B50090B757 /* test.mov in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};

SJSAssetExportSession/ExportSession.swift

@@ -25,11 +25,11 @@ public final class ExportSession {
         }
     }
 
-    public func export(
+    public static func export(
         asset: sending AVAsset,
         audioMix: sending AVAudioMix?,
         audioOutputSettings: [String: (any Sendable)],
-        videoComposition: sending AVVideoComposition?,
+        videoComposition: sending AVVideoComposition,
         videoOutputSettings: [String: (any Sendable)],
         timeRange: CMTimeRange? = nil,
         optimizeForNetworkUse: Bool = false,

SJSAssetExportSession/SampleWriter.swift

@@ -6,6 +6,9 @@
 //
 
 import AVFoundation.AVAsset
+import OSLog
+
+private let log = Logger(subsystem: "SJSAssetExportSession", category: "SampleWriter")
 
 private extension AVAsset {
     func sendTracks(withMediaType mediaType: AVMediaType) async throws -> sending [AVAssetTrack] {
@@ -24,25 +27,21 @@ actor SampleWriter {
         queue.asUnownedSerialExecutor()
     }
 
-    let audioTracks: [AVAssetTrack]
+    private let audioMix: AVAudioMix?
 
-    let audioMix: AVAudioMix?
+    private let audioOutputSettings: [String: (any Sendable)]
 
-    let audioOutputSettings: [String: (any Sendable)]
+    private let videoComposition: AVVideoComposition?
 
-    let videoTracks: [AVAssetTrack]
+    private let videoOutputSettings: [String: (any Sendable)]
 
-    let videoComposition: AVVideoComposition?
+    private let reader: AVAssetReader
 
-    let videoOutputSettings: [String: (any Sendable)]
+    private let writer: AVAssetWriter
 
-    let reader: AVAssetReader
+    private let duration: CMTime
 
-    let writer: AVAssetWriter
+    private let timeRange: CMTimeRange
-
-    let duration: CMTime
-
-    let timeRange: CMTimeRange
 
     private var audioOutput: AVAssetReaderAudioMixOutput?
 
@@ -52,14 +51,18 @@ actor SampleWriter {
 
     private var videoInput: AVAssetWriterInput?
 
-    private var pixelBufferAdaptor: AVAssetWriterInputPixelBufferAdaptor?
+    private lazy var progressStream: AsyncStream<Float> = AsyncStream { continuation in
+        progressContinuation = continuation
+    }
+
+    private var progressContinuation: AsyncStream<Float>.Continuation?
 
     init(
         asset: sending AVAsset,
         timeRange: CMTimeRange,
         audioMix: AVAudioMix?,
         audioOutputSettings: sending [String: (any Sendable)],
-        videoComposition: AVVideoComposition?,
+        videoComposition: AVVideoComposition,
         videoOutputSettings: sending [String: (any Sendable)],
         optimizeForNetworkUse: Bool,
         outputURL: URL,
@@ -81,26 +84,31 @@ actor SampleWriter {
             throw ExportSession.Error.setupFailure(reason: "Cannot apply video output settings")
         }
 
-        self.audioTracks = try await asset.sendTracks(withMediaType: .audio)
+        let audioTracks = try await asset.sendTracks(withMediaType: .audio)
+        let videoTracks = try await asset.sendTracks(withMediaType: .video)
+
         self.audioMix = audioMix
         self.audioOutputSettings = audioOutputSettings
-        self.videoTracks = try await asset.sendTracks(withMediaType: .video)
         self.videoComposition = videoComposition
         self.videoOutputSettings = videoOutputSettings
         self.reader = reader
         self.writer = writer
         self.duration = duration
         self.timeRange = timeRange
+
+        try setUpAudio(audioTracks: audioTracks)
+        try setUpVideo(videoTracks: videoTracks)
     }
 
     func writeSamples() async throws {
+        progressContinuation?.yield(0)
+
         writer.startWriting()
         reader.startReading()
         writer.startSession(atSourceTime: timeRange.start)
 
-        async let audioResult = try encodeAudioTracks(audioTracks)
+        await encodeAudioTracks()
-        async let videoResult = try encodeVideoTracks(videoTracks)
+        await encodeVideoTracks()
-        _ = try await (audioResult, videoResult)
 
         if reader.status == .cancelled || writer.status == .cancelled {
             throw CancellationError()
@@ -119,10 +127,12 @@ actor SampleWriter {
         }
     }
 
-    private func encodeAudioTracks(_ audioTracks: [AVAssetTrack]) async throws -> Bool {
+    private func setUpAudio(audioTracks: [AVAssetTrack]) throws {
-        guard !audioTracks.isEmpty else { return false }
+        guard !audioTracks.isEmpty else { return }
 
         let audioOutput = AVAssetReaderAudioMixOutput(audioTracks: audioTracks, audioSettings: nil)
+        audioOutput.alwaysCopiesSampleData = false
+        audioOutput.audioMix = audioMix
         guard reader.canAdd(audioOutput) else {
             throw ExportSession.Error.setupFailure(reason: "Can't add audio output to reader")
         }
@@ -130,55 +140,23 @@ actor SampleWriter {
         self.audioOutput = audioOutput
 
         let audioInput = AVAssetWriterInput(mediaType: .audio, outputSettings: audioOutputSettings)
+        audioInput.expectsMediaDataInRealTime = false
         guard writer.canAdd(audioInput) else {
             throw ExportSession.Error.setupFailure(reason: "Can't add audio input to writer")
         }
         writer.add(audioInput)
         self.audioInput = audioInput
-
-        return await withCheckedContinuation { continuation in
-            self.audioInput?.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { $0.writeReadyAudioSamples() }
-                if !hasMoreSamples {
-                    continuation.resume(returning: true)
-                }
-            }
-        }
     }
 
-    private func writeReadyAudioSamples() -> Bool {
+    private func setUpVideo(videoTracks: [AVAssetTrack]) throws {
-        guard let audioOutput, let audioInput else { return true }
+        guard !videoTracks.isEmpty else {
-
+            throw ExportSession.Error.setupFailure(reason: "No video tracks")
-        while audioInput.isReadyForMoreMediaData {
-            guard reader.status == .reading && writer.status == .writing,
-                  let sampleBuffer = audioOutput.copyNextSampleBuffer() else {
-                audioInput.markAsFinished()
-                NSLog("Finished encoding ready audio samples from \(audioOutput)")
-                return false
-            }
-
-            guard audioInput.append(sampleBuffer) else {
-                NSLog("Failed to append audio sample buffer \(sampleBuffer) to input \(audioInput)")
-                return false
-            }
         }
 
-        // Everything was appended successfully, return true indicating there's more to do.
+        let videoOutput = AVAssetReaderVideoCompositionOutput(
-        NSLog("Completed encoding ready audio samples, more to come...")
+            videoTracks: videoTracks,
-        return true
+            videoSettings: nil
-    }
+        )
-
-    private func encodeVideoTracks(_ videoTracks: [AVAssetTrack]) async throws -> Bool {
-        guard !videoTracks.isEmpty else { return false }
-
-        guard let width = videoComposition.map({ Int($0.renderSize.width) })
-                ?? (videoOutputSettings[AVVideoWidthKey] as? NSNumber)?.intValue,
-              let height = videoComposition.map({ Int($0.renderSize.height) })
-                ?? (videoOutputSettings[AVVideoHeightKey] as? NSNumber)?.intValue else {
-            throw ExportSession.Error.setupFailure(reason: "Export dimensions must be provided in a video composition or video output settings")
-        }
-
-        let videoOutput = AVAssetReaderVideoCompositionOutput(videoTracks: videoTracks, videoSettings: nil)
         videoOutput.alwaysCopiesSampleData = false
         videoOutput.videoComposition = videoComposition
         guard reader.canAdd(videoOutput) else {
@@ -188,72 +166,60 @@ actor SampleWriter {
         self.videoOutput = videoOutput
 
         let videoInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoOutputSettings)
+        videoInput.expectsMediaDataInRealTime = false
         guard writer.canAdd(videoInput) else {
             throw ExportSession.Error.setupFailure(reason: "Can't add video input to writer")
         }
         writer.add(videoInput)
         self.videoInput = videoInput
-
-        let pixelBufferAttributes: [String: Any] = [
-            kCVPixelBufferPixelFormatTypeKey as String: NSNumber(integerLiteral: Int(kCVPixelFormatType_32RGBA)),
-            kCVPixelBufferWidthKey as String: NSNumber(integerLiteral: width),
-            kCVPixelBufferHeightKey as String: NSNumber(integerLiteral: height),
-            "IOSurfaceOpenGLESTextureCompatibility": NSNumber(booleanLiteral: true),
-            "IOSurfaceOpenGLESFBOCompatibility": NSNumber(booleanLiteral: true),
-        ]
-        pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(
-            assetWriterInput: videoInput,
-            sourcePixelBufferAttributes: pixelBufferAttributes
-        )
-
-        return await withCheckedContinuation { continuation in
-            self.videoInput?.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { $0.writeReadyVideoSamples() }
-                if !hasMoreSamples {
-                    continuation.resume(returning: true)
-                }
-            }
-        }
     }
 
-    private func writeReadyVideoSamples() -> Bool {
+    private func writeReadySamples(output: AVAssetReaderOutput, input: AVAssetWriterInput) -> Bool {
-        guard let videoOutput, let videoInput, let pixelBufferAdaptor else { return true }
+        while input.isReadyForMoreMediaData {
-
-        while videoInput.isReadyForMoreMediaData {
             guard reader.status == .reading && writer.status == .writing,
-                  let sampleBuffer = videoOutput.copyNextSampleBuffer() else {
+                  let sampleBuffer = output.copyNextSampleBuffer() else {
-                videoInput.markAsFinished()
+                input.markAsFinished()
-                NSLog("Finished encoding ready video samples from \(videoOutput)")
+                log.debug("Finished encoding ready audio samples from \(output)")
                 return false
             }
 
-            let samplePresentationTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer) - timeRange.start
+            guard input.append(sampleBuffer) else {
-            let progress = Float(samplePresentationTime.seconds / duration.seconds)
+                log.error("""
-#warning("TODO: publish progress to an AsyncStream")
+                    Failed to append audio sample buffer \(String(describing: sampleBuffer)) to
-
+                    input \(input.debugDescription)
-            guard let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool else {
+                """)
-                NSLog("No pixel buffer pool available on adaptor \(pixelBufferAdaptor)")
                 return false
             }
-            var toRenderBuffer: CVPixelBuffer?
-            let result = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, pixelBufferPool, &toRenderBuffer)
-            var handled = false
-            if result == kCVReturnSuccess, let toBuffer = toRenderBuffer {
-                handled = pixelBufferAdaptor.append(toBuffer, withPresentationTime: samplePresentationTime)
-                if !handled { return false }
-            }
-            if !handled {
-#warning("is this really necessary?! seems like a failure scenario...")
-                guard videoInput.append(sampleBuffer) else {
-                    NSLog("Failed to append video sample buffer \(sampleBuffer) to input \(videoInput)")
-                    return false
-                }
-            }
         }
 
         // Everything was appended successfully, return true indicating there's more to do.
-        NSLog("Completed encoding ready video samples, more to come...")
+        log.debug("Completed encoding ready audio samples, more to come...")
         return true
     }
 
+    private func encodeAudioTracks() async {
+        return await withCheckedContinuation { continuation in
+            self.audioInput!.requestMediaDataWhenReady(on: queue) {
+                let hasMoreSamples = self.assumeIsolated { _self in
+                    _self.writeReadySamples(output: _self.audioOutput!, input: _self.audioInput!)
+                }
+                if !hasMoreSamples {
+                    continuation.resume()
+                }
+            }
+        }
+    }
+
+    private func encodeVideoTracks() async {
+        return await withCheckedContinuation { continuation in
+            self.videoInput!.requestMediaDataWhenReady(on: queue) {
+                let hasMoreSamples = self.assumeIsolated { _self in
+                    _self.writeReadySamples(output: _self.videoOutput!, input: _self.videoInput!)
+                }
+                if !hasMoreSamples {
+                    continuation.resume()
+                }
+            }
+        }
+    }
 }

SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift

@@ -5,13 +5,58 @@
 //  Created by Sami Samhuri on 2024-06-29.
 //
 
-import Testing
+import AVFoundation
 @testable import SJSAssetExportSession
+import Testing
 
-struct SJSAssetExportSessionTests {
+final class ExportSessionTests {
+    @Test func test_encode_h264_720p_30fps() async throws {
+        let sourceURL = Bundle(for: Self.self).url(forResource: "test", withExtension: "mov")!
+        let sourceAsset = AVURLAsset(url: sourceURL)
+        let timestamp = Int(Date.now.timeIntervalSince1970)
+        let filename = "ExportSessionTests_testEncode_\(timestamp).mp4"
+        let destinationURL = URL.temporaryDirectory.appending(component: filename)
+        defer { _ = try? FileManager.default.removeItem(at: destinationURL) }
 
-    @Test func testExample() async throws {
+        let size = CGSize(width: 1280, height: 720)
-        // Write your test here and use APIs like `#expect(...)` to check expected conditions.
+        let duration = CMTime(seconds: 1, preferredTimescale: 600)
+        let videoComposition = try await AVMutableVideoComposition.videoComposition(withPropertiesOf: sourceAsset)
+        videoComposition.renderSize = size
+        videoComposition.renderScale = 1
+        videoComposition.frameDuration = CMTime(value: 1, timescale: 30)
+        try await ExportSession.export(
+            asset: sourceAsset,
+            audioMix: nil,
+            audioOutputSettings: [
+                AVFormatIDKey: kAudioFormatMPEG4AAC,
+                AVNumberOfChannelsKey: NSNumber(value: 2),
+                AVSampleRateKey: NSNumber(value: 44_100.0),
+            ],
+            videoComposition: videoComposition,
+            videoOutputSettings: [
+                AVVideoCodecKey: AVVideoCodecType.h264.rawValue,
+                AVVideoWidthKey: NSNumber(value: Int(size.width)),
+                AVVideoHeightKey: NSNumber(value: Int(size.height)),
+                AVVideoCompressionPropertiesKey: [
+                    AVVideoAverageBitRateKey: NSNumber(value: 1_000_000),
+                    AVVideoProfileLevelKey: AVVideoProfileLevelH264HighAutoLevel as String,
+                ] as [String: any Sendable],
+                AVVideoColorPropertiesKey: [
+                    AVVideoColorPrimariesKey: AVVideoColorPrimaries_ITU_R_709_2,
+                    AVVideoTransferFunctionKey: AVVideoTransferFunction_ITU_R_709_2,
+                    AVVideoYCbCrMatrixKey: AVVideoYCbCrMatrix_ITU_R_709_2,
+                ],
+            ],
+            timeRange: CMTimeRange(start: .zero, duration: duration),
+            to: destinationURL,
+            as: .mp4
+        )
+
+        let asset = AVURLAsset(url: destinationURL)
+        #expect(try await asset.load(.duration) == duration)
+        try #require(await asset.loadTracks(withMediaType: .video).count == 1)
+        let videoTrack = try #require(try await asset.loadTracks(withMediaType: .video).first)
+        #expect(try await videoTrack.load(.naturalSize) == size)
+        try #require(try await asset.loadTracks(withMediaType: .audio).count == 1)
     }
-
 }