from __future__ import annotations

import unittest

from llm_stream import OpenAIStreamAccumulator


class OpenAIStreamAccumulatorTests(unittest.TestCase):
    def test_repeated_delta_chunks_are_preserved_as_model_output(self) -> None:
        accumulator = OpenAIStreamAccumulator()
        self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes")
        self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes")
        self.assertEqual(accumulator.text, "yesyes")

    def test_final_snapshot_can_add_a_missing_suffix(self) -> None:
        accumulator = OpenAIStreamAccumulator()
        self.assertEqual(accumulator.feed({"delta": {"content": "first"}}), "first")
        self.assertEqual(
            accumulator.feed({"message": {"content": "first second"}}),
            " second",
        )
        self.assertEqual(accumulator.text, "first second")

    def test_incompatible_final_snapshot_is_not_appended_twice(self) -> None:
        accumulator = OpenAIStreamAccumulator()
        accumulator.feed({"delta": {"content": "streamed answer"}})
        self.assertEqual(
            accumulator.feed({"message": {"content": "rewritten final answer"}}),
            "",
        )
        self.assertEqual(accumulator.text, "streamed answer")


if __name__ == "__main__":
    unittest.main()