def test_streaming_citation_handling(self):
        """Test citation handling with simulated streaming updates."""
        tracker = CitationTracker()

        # Simulate a streaming scenario where text comes in chunks
        chunks = [
            "This is the first chunk ",
            "with no citations. This is the second chunk with a ",
            "citation [abc1234] and some more text. ",
            "This is the third chunk with another citation [def5678] ",
            "and the first citation again [abc1234] in a new position."
        ]

        all_text = ""
        total_spans_found = 0

        for i, chunk in enumerate(chunks):
            chunk_start = len(all_text)
            all_text += chunk

            # For streaming, we need to extract citation spans from the chunk
            # and check if they are new in the context of the accumulated text
            pattern = r'\[([\w]{7,8})\]'
            for match in re.finditer(pattern, chunk):
                citation_id = match.group(1)
                start = match.start() + chunk_start
                end = match.end() + chunk_start

                # Check if this span is new for this citation ID
                if tracker.is_new_span(citation_id, (start, end)):
                    total_spans_found += 1

        # Check final state
        assert "abc1234" in tracker.processed_spans, "First citation should be tracked"
        assert "def5678" in tracker.processed_spans, "Second citation should be tracked"
        assert len(tracker.processed_spans["abc1234"]) == 2, "First citation should have 2 spans"
        assert len(tracker.processed_spans["def5678"]) == 1, "Second citation should have 1 span"
        assert total_spans_found == 3, "Should have found 3 spans in total"
Test citation handling with simulated streaming updates.
py/tests/unit/retrieval/test_citations.py