import asyncio
import sys
import os

# Add src to path
sys.path.append(os.path.join(os.getcwd(), "src"))

async def test_streaming_fix():
    print("Starting standalone verification for StopIteration fix...")
    
    try:
        from ai_med_extract.utils.unified_model_manager import TransformersModel, GenerationConfig
        
        class MockTransformersModel(TransformersModel):
            def __init__(self):
                self.name = "mock"
                self.model_type = "summarization"
                self._status = "loaded" 
                self._model = object()
                self._model_obj = object()
                self._tokenizer = object()

            def generate_stream(self, prompt, config):
                yield "Token 1"
                yield "Token 2"
                yield "Token 3"
        
        print("Testing MockTransformersModel...")
        mock = MockTransformersModel()
        tokens = []
        
        # Test generate_async_stream
        async for token in mock.generate_async_stream("test prompt", GenerationConfig()):
            print(f"Received: {token}")
            tokens.append(token)
        
        if tokens == ["Token 1", "Token 2", "Token 3"]:
            print("SUCCESS: MockTransformersModel verification successful! No StopIteration raised.")
        else:
            print(f"FAILURE: Expected ['Token 1', 'Token 2', 'Token 3'], got {tokens}")
            sys.exit(1)
        
    except Exception as e:
        print(f"ERROR: Verification failed with exception: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

    print("\nVerification complete.")

if __name__ == "__main__":
    asyncio.run(test_streaming_fix())