Repository: vapor/vapor
File: Tests/VaporTests/RouteTests.swift (Swift)

Code before:
        app.on(.POST, "1gb", body: .collect(maxSize: "1gb")) { request in
            HTTPStatus.ok
        }

        var buffer = ByteBufferAllocator().buffer(capacity: 0)
        buffer.writeBytes(Array(repeating: 0, count: 500_000))
        try app.testable(method: .running).test(.POST, "/default", body: buffer) { res in
            XCTAssertEqual(res.status, .payloadTooLarge)
        }.test(.POST, "/1kb", body: buffer) { res in
            XCTAssertEqual(res.status, .payloadTooLarge)
        }.test(.POST, "/1mb", body: buffer) { res in
            XCTAssertEqual(res.status, .ok)
        }.test(.POST, "/1gb", body: buffer) { res in
            XCTAssertEqual(res.status, .ok)
        }
    }
    
    func testWebsocketUpgrade() throws {
        let app = Application(.testing)
        defer { app.shutdown() }
        
        let testMarkerHeaderKey = "TestMarker"
        let testMarkerHeaderValue = "addedInShouldUpgrade"
        
        app.routes.webSocket("customshouldupgrade", shouldUpgrade: { req in
            return req.eventLoop.future([testMarkerHeaderKey:testMarkerHeaderValue])
        }, onUpgrade: { _, _ in })
        
        try app.testable(method: .running).test(.GET, "customshouldupgrade", beforeRequest: { req in
            req.headers.replaceOrAdd(name: HTTPHeaders.Name.secWebSocketVersion, value: "13")
            req.headers.replaceOrAdd(name: HTTPHeaders.Name.secWebSocketKey, value: "zyFJtLIpI2ASsmMHJ4Cf0A==")
            req.headers.replaceOrAdd(name: .connection, value: "Upgrade")
            req.headers.replaceOrAdd(name: .upgrade, value: "websocket")
        }) { res in
            XCTAssertEqual(res.headers.first(name: testMarkerHeaderKey), testMarkerHeaderValue)
        }
    }
}


Review comment (suggestion):
```suggestion
            return req.eventLoop.future([testMarkerHeaderKey: testMarkerHeaderValue])
```

Code after:
        app.on(.POST, "1gb", body: .collect(maxSize: "1gb")) { request in
            HTTPStatus.ok
        }

        var buffer = ByteBufferAllocator().buffer(capacity: 0)
        buffer.writeBytes(Array(repeating: 0, count: 500_000))
        try app.testable(method: .running).test(.POST, "/default", body: buffer) { res in
            XCTAssertEqual(res.status, .payloadTooLarge)
        }.test(.POST, "/1kb", body: buffer) { res in
            XCTAssertEqual(res.status, .payloadTooLarge)
        }.test(.POST, "/1mb", body: buffer) { res in
            XCTAssertEqual(res.status, .ok)
        }.test(.POST, "/1gb", body: buffer) { res in
            XCTAssertEqual(res.status, .ok)
        }
    }
    
    func testWebsocketUpgrade() throws {
        let app = Application(.testing)
        defer { app.shutdown() }
        
        let testMarkerHeaderKey = "TestMarker"
        let testMarkerHeaderValue = "addedInShouldUpgrade"
        
        app.routes.webSocket("customshouldupgrade", shouldUpgrade: { req in
            return req.eventLoop.future([testMarkerHeaderKey: testMarkerHeaderValue])
        }, onUpgrade: { _, _ in })
        
        try app.testable(method: .running).test(.GET, "customshouldupgrade", beforeRequest: { req in
            req.headers.replaceOrAdd(name: HTTPHeaders.Name.secWebSocketVersion, value: "13")
            req.headers.replaceOrAdd(name: HTTPHeaders.Name.secWebSocketKey, value: "zyFJtLIpI2ASsmMHJ4Cf0A==")
            req.headers.replaceOrAdd(name: .connection, value: "Upgrade")
            req.headers.replaceOrAdd(name: .upgrade, value: "websocket")
        }) { res in
            XCTAssertEqual(res.headers.first(name: testMarkerHeaderKey), testMarkerHeaderValue)
        }
    }
}