add e2e test for simple streaming server (#247)

* test: add e2e test for simple streaming server

* test: remove redundant assertion and simplify output checks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: AdolfoVillalobos

tests/e2e/default_single_streaming.py

@@ -0,0 +1,23 @@
+from litserve import LitServer, LitAPI
+
+
+class SimpleStreamingAPI(LitAPI):
+    def setup(self, device) -> None:
+        self.model = lambda x, y: x * y
+
+    def decode_request(self, request):
+        return request["input"]
+
+    def predict(self, x):
+        for i in range(1, 4):
+            yield self.model(i, x)
+
+    def encode_response(self, output_stream):
+        for output in output_stream:
+            yield {"output": output}
+
+
+if __name__ == "__main__":
+    api = SimpleStreamingAPI()
+    server = LitServer(api, stream=True)
+    server.run(port=8000)

tests/e2e/test_e2e.py

@@ -295,3 +295,21 @@ def test_openai_parity_with_response_format():
         assert r.choices[0].delta.content == expected_out, (
             f"Server didn't return expected output.\n" f"OpenAI client output: {r}"
         )
+
+
+@e2e_from_file("tests/e2e/default_single_streaming.py")
+def test_e2e_single_streaming():
+    resp = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0}, headers=None, stream=True)
+    assert resp.status_code == 200, f"Expected response to be 200 but got {resp.status_code}"
+
+    outputs = []
+    for line in resp.iter_lines():
+        if line:
+            outputs.append(json.loads(line.decode("utf-8")))
+
+    assert len(outputs) == 3, "Expected 3 streamed outputs"
+    assert outputs[-1] == {"output": 12.0}, "Final output doesn't match expected value"
+
+    expected_values = [4.0, 8.0, 12.0]
+    for i, output in enumerate(outputs):
+        assert output["output"] == expected_values[i], f"Intermediate output {i} is not expected value"