@@ -574,6 +574,7 @@ struct LambdaResponseStreamWriterHeadersTests {
574
574
final class MockLambdaResponseStreamWriter : LambdaResponseStreamWriter {
575
575
private( set) var writtenBuffers : [ ByteBuffer ] = [ ]
576
576
private( set) var isFinished = false
577
+ private( set) var hasCustomHeaders = false
577
578
578
579
func write( _ buffer: ByteBuffer ) async throws {
579
580
writtenBuffers. append ( buffer)
@@ -589,9 +590,7 @@ final class MockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
589
590
}
590
591
591
592
func writeCustomHeader( _ buffer: NIOCore . ByteBuffer ) async throws {
592
- // This is a mock, so we don't actually write custom headers.
593
- // In a real implementation, this would handle writing custom headers.
594
- fatalError ( " Unexpected call to writeCustomHeader " )
593
+ hasCustomHeaders = true
595
594
}
596
595
}
597
596
@@ -602,6 +601,7 @@ final class FailingMockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
602
601
private( set) var writtenBuffers : [ ByteBuffer ] = [ ]
603
602
private( set) var writeCallCount = 0
604
603
private( set) var isFinished = false
604
+ private( set) var hasCustomHeaders = false
605
605
private let failOnWriteCall : Int
606
606
607
607
init ( failOnWriteCall: Int ) {
@@ -628,9 +628,7 @@ final class FailingMockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
628
628
}
629
629
630
630
func writeCustomHeader( _ buffer: NIOCore . ByteBuffer ) async throws {
631
- // This is a mock, so we don't actually write custom headers.
632
- // In a real implementation, this would handle writing custom headers.
633
- fatalError ( " Unexpected call to writeCustomHeader " )
631
+ hasCustomHeaders = true
634
632
}
635
633
}
636
634
@@ -712,6 +710,7 @@ final class TrackingLambdaResponseStreamWriter: LambdaResponseStreamWriter {
712
710
private( set) var finishCallCount = 0
713
711
private( set) var writeAndFinishCallCount = 0
714
712
private( set) var isFinished = false
713
+ private( set) var hasCustomHeaders = false
715
714
716
715
func write( _ buffer: ByteBuffer ) async throws {
717
716
writeCallCount += 1
@@ -730,9 +729,7 @@ final class TrackingLambdaResponseStreamWriter: LambdaResponseStreamWriter {
730
729
}
731
730
732
731
func writeCustomHeader( _ buffer: NIOCore . ByteBuffer ) async throws {
733
- // This is a mock, so we don't actually write custom headers.
734
- // In a real implementation, this would handle writing custom headers.
735
- fatalError ( " Unexpected call to writeCustomHeader " )
732
+ hasCustomHeaders = true
736
733
}
737
734
}
738
735
@@ -741,6 +738,7 @@ final class CustomBehaviorLambdaResponseStreamWriter: LambdaResponseStreamWriter
741
738
private( set) var writtenBuffers : [ ByteBuffer ] = [ ]
742
739
private( set) var customBehaviorTriggered = false
743
740
private( set) var isFinished = false
741
+ private( set) var hasCustomHeaders = false
744
742
745
743
func write( _ buffer: ByteBuffer ) async throws {
746
744
// Trigger custom behavior on any write
@@ -759,8 +757,6 @@ final class CustomBehaviorLambdaResponseStreamWriter: LambdaResponseStreamWriter
759
757
}
760
758
761
759
func writeCustomHeader( _ buffer: NIOCore . ByteBuffer ) async throws {
762
- // This is a mock, so we don't actually write custom headers.
763
- // In a real implementation, this would handle writing custom headers.
764
- fatalError ( " Unexpected call to writeCustomHeader " )
760
+ hasCustomHeaders = true
765
761
}
766
762
}
0 commit comments