|
1 | 1 | import asyncio |
| 2 | +from unittest.mock import patch |
2 | 3 |
|
3 | 4 | from cubepi.providers.base import ( |
4 | 5 | Model, |
|
9 | 10 | ) |
10 | 11 | from cubepi.providers.faux import ( |
11 | 12 | FauxProvider, |
| 13 | + _can_accept_extended_args, |
12 | 14 | faux_assistant_message, |
13 | 15 | faux_text, |
14 | 16 | faux_thinking, |
@@ -639,3 +641,300 @@ async def test_prompt_cache_property_is_copy(self): |
639 | 641 | cache = provider.prompt_cache |
640 | 642 | cache["injected"] = "value" |
641 | 643 | assert "injected" not in provider.prompt_cache |
| 644 | + |
| 645 | + |
| 646 | +class TestCanAcceptExtendedArgs: |
| 647 | + """Tests for _can_accept_extended_args edge cases (lines 147-148).""" |
| 648 | + |
| 649 | + def test_returns_false_when_signature_raises_value_error(self): |
| 650 | + """When inspect.signature raises ValueError, should return False.""" |
| 651 | + with patch( |
| 652 | + "cubepi.providers.faux.inspect.signature", side_effect=ValueError("boom") |
| 653 | + ): |
| 654 | + assert _can_accept_extended_args(lambda: None) is False |
| 655 | + |
| 656 | + def test_returns_false_when_signature_raises_type_error(self): |
| 657 | + """When inspect.signature raises TypeError, should return False.""" |
| 658 | + with patch( |
| 659 | + "cubepi.providers.faux.inspect.signature", side_effect=TypeError("boom") |
| 660 | + ): |
| 661 | + assert _can_accept_extended_args(lambda: None) is False |
| 662 | + |
| 663 | + |
| 664 | +class TestFauxProviderProduceExceptionHandling: |
| 665 | + """Tests for the _produce exception handler (lines 279-290).""" |
| 666 | + |
| 667 | + def _make_model(self): |
| 668 | + return Model(id="faux-1", provider="faux") |
| 669 | + |
| 670 | + async def test_factory_raising_exception_produces_error_result(self): |
| 671 | + """When a factory raises an Exception, _produce catches it and |
| 672 | + produces an error AssistantMessage.""" |
| 673 | + |
| 674 | + def bad_factory(messages, model): |
| 675 | + raise RuntimeError("factory exploded") |
| 676 | + |
| 677 | + provider = FauxProvider() |
| 678 | + provider.set_responses([bad_factory]) |
| 679 | + model = self._make_model() |
| 680 | + |
| 681 | + stream = await provider.stream(model, []) |
| 682 | + events = [e async for e in stream] |
| 683 | + result = await stream.result() |
| 684 | + |
| 685 | + assert result.stop_reason == "error" |
| 686 | + assert "factory exploded" in (result.error_message or "") |
| 687 | + assert any(e.type == "error" for e in events) |
| 688 | + |
| 689 | + async def test_factory_raising_base_exception_reraises(self): |
| 690 | + """When a factory raises a BaseException (not Exception), |
| 691 | + _produce catches it, sets an error result, and re-raises. |
| 692 | + The re-raised exception surfaces as the task exception.""" |
| 693 | + |
| 694 | + class CustomBaseException(BaseException): |
| 695 | + pass |
| 696 | + |
| 697 | + def bad_factory(messages, model): |
| 698 | + raise CustomBaseException("base boom") |
| 699 | + |
| 700 | + provider = FauxProvider() |
| 701 | + provider.set_responses([bad_factory]) |
| 702 | + model = self._make_model() |
| 703 | + |
| 704 | + stream = await provider.stream(model, []) |
| 705 | + _ = [e async for e in stream] |
| 706 | + result = await stream.result() |
| 707 | + |
| 708 | + # The BaseException path still sets an error result |
| 709 | + assert result.stop_reason == "error" |
| 710 | + assert "base boom" in (result.error_message or "") |
| 711 | + |
| 712 | + |
| 713 | +class TestFauxProviderAbortDuringBlocks: |
| 714 | + """Tests for abort signal checks during block iteration and chunk streaming.""" |
| 715 | + |
| 716 | + def _make_model(self): |
| 717 | + return Model(id="faux-1", provider="faux") |
| 718 | + |
| 719 | + async def test_abort_between_blocks(self): |
| 720 | + """Abort signal set between blocks triggers the block-level abort |
| 721 | + check (lines 318-325). |
| 722 | +
|
| 723 | + Strategy: directly call _stream_with_deltas with a pre-constructed |
| 724 | + message and set the signal synchronously during the 'thinking_end' |
| 725 | + push via a patched MessageStream.push, so it's set before the |
| 726 | + for-block check runs for the next block. |
| 727 | + """ |
| 728 | + from cubepi.providers.base import MessageStream |
| 729 | + |
| 730 | + provider = FauxProvider(token_size_min=100, token_size_max=100) |
| 731 | + signal = asyncio.Event() |
| 732 | + |
| 733 | + message = faux_assistant_message([faux_thinking("ok"), faux_text("answer")]) |
| 734 | + |
| 735 | + ms = MessageStream() |
| 736 | + |
| 737 | + original_push = ms.push |
| 738 | + |
| 739 | + def push_and_set_signal(event): |
| 740 | + original_push(event) |
| 741 | + if event.type == "thinking_end": |
| 742 | + signal.set() |
| 743 | + |
| 744 | + ms.push = push_and_set_signal # type: ignore[assignment] |
| 745 | + |
| 746 | + await provider._stream_with_deltas(ms, message, signal) |
| 747 | + |
| 748 | + result = await ms.result() |
| 749 | + assert result.stop_reason == "aborted" |
| 750 | + |
| 751 | + async def test_abort_during_thinking_chunks(self): |
| 752 | + """Abort signal set while thinking deltas are being streamed |
| 753 | + (lines 340-347).""" |
| 754 | + # Use long thinking text to ensure multiple chunks |
| 755 | + long_thinking = "a" * 200 |
| 756 | + provider = FauxProvider(token_size_min=1, token_size_max=1) |
| 757 | + provider.set_responses( |
| 758 | + [ |
| 759 | + faux_assistant_message( |
| 760 | + [faux_thinking(long_thinking), faux_text("answer")] |
| 761 | + ) |
| 762 | + ] |
| 763 | + ) |
| 764 | + model = self._make_model() |
| 765 | + signal = asyncio.Event() |
| 766 | + |
| 767 | + stream = await provider.stream(model, [], options=StreamOptions(signal=signal)) |
| 768 | + |
| 769 | + events = [] |
| 770 | + thinking_delta_count = 0 |
| 771 | + async for event in stream: |
| 772 | + events.append(event) |
| 773 | + if event.type == "thinking_delta": |
| 774 | + thinking_delta_count += 1 |
| 775 | + # Abort after a few thinking deltas |
| 776 | + if thinking_delta_count >= 3: |
| 777 | + signal.set() |
| 778 | + |
| 779 | + result = await stream.result() |
| 780 | + assert result.stop_reason == "aborted" |
| 781 | + assert any(e.type == "error" for e in events) |
| 782 | + # We should have some thinking deltas but not all of them |
| 783 | + assert thinking_delta_count >= 3 |
| 784 | + # The text block should NOT have started |
| 785 | + event_types = [e.type for e in events] |
| 786 | + assert "text_start" not in event_types |
| 787 | + |
| 788 | + async def test_abort_during_tool_call_chunks(self): |
| 789 | + """Abort signal set while tool call deltas are being streamed |
| 790 | + (lines 425-432).""" |
| 791 | + # Use a large arguments dict to produce multiple chunks |
| 792 | + large_args = {f"key_{i}": f"value_{i}" for i in range(20)} |
| 793 | + provider = FauxProvider(token_size_min=1, token_size_max=1) |
| 794 | + provider.set_responses( |
| 795 | + [ |
| 796 | + faux_assistant_message( |
| 797 | + [faux_tool_call("search", large_args, id="tc-1")], |
| 798 | + stop_reason="tool_use", |
| 799 | + ) |
| 800 | + ] |
| 801 | + ) |
| 802 | + model = self._make_model() |
| 803 | + signal = asyncio.Event() |
| 804 | + |
| 805 | + stream = await provider.stream(model, [], options=StreamOptions(signal=signal)) |
| 806 | + |
| 807 | + events = [] |
| 808 | + toolcall_delta_count = 0 |
| 809 | + async for event in stream: |
| 810 | + events.append(event) |
| 811 | + if event.type == "toolcall_delta": |
| 812 | + toolcall_delta_count += 1 |
| 813 | + # Abort after a few tool call deltas |
| 814 | + if toolcall_delta_count >= 3: |
| 815 | + signal.set() |
| 816 | + |
| 817 | + result = await stream.result() |
| 818 | + assert result.stop_reason == "aborted" |
| 819 | + assert any(e.type == "error" for e in events) |
| 820 | + assert toolcall_delta_count >= 3 |
| 821 | + # Tool call should NOT have ended normally |
| 822 | + event_types = [e.type for e in events] |
| 823 | + assert "toolcall_end" not in event_types |
| 824 | + |
| 825 | + async def test_abort_during_text_then_tool_blocks(self): |
| 826 | + """Abort during text block prevents tool call block from starting.""" |
| 827 | + long_text = "word " * 100 |
| 828 | + provider = FauxProvider(token_size_min=1, token_size_max=1) |
| 829 | + provider.set_responses( |
| 830 | + [ |
| 831 | + faux_assistant_message( |
| 832 | + [ |
| 833 | + faux_text(long_text), |
| 834 | + faux_tool_call("search", {"q": "test"}, id="tc-1"), |
| 835 | + ], |
| 836 | + stop_reason="tool_use", |
| 837 | + ) |
| 838 | + ] |
| 839 | + ) |
| 840 | + model = self._make_model() |
| 841 | + signal = asyncio.Event() |
| 842 | + |
| 843 | + stream = await provider.stream(model, [], options=StreamOptions(signal=signal)) |
| 844 | + |
| 845 | + events = [] |
| 846 | + text_delta_count = 0 |
| 847 | + async for event in stream: |
| 848 | + events.append(event) |
| 849 | + if event.type == "text_delta": |
| 850 | + text_delta_count += 1 |
| 851 | + if text_delta_count >= 3: |
| 852 | + signal.set() |
| 853 | + |
| 854 | + result = await stream.result() |
| 855 | + assert result.stop_reason == "aborted" |
| 856 | + # Tool call block should never start |
| 857 | + event_types = [e.type for e in events] |
| 858 | + assert "toolcall_start" not in event_types |
| 859 | + |
| 860 | + |
| 861 | +class TestFauxProviderCacheTokenCalculation: |
| 862 | + """Tests for cache token calculation logic (lines 185-216).""" |
| 863 | + |
| 864 | + def _make_model(self): |
| 865 | + return Model(id="faux-1", provider="faux") |
| 866 | + |
| 867 | + async def test_cache_usage_first_call_structure(self): |
| 868 | + """First call: input_tokens == prompt_tokens, cache_write == prompt_tokens, |
| 869 | + cache_read == 0.""" |
| 870 | + provider = FauxProvider() |
| 871 | + provider.set_responses([faux_assistant_message("hello")]) |
| 872 | + model = self._make_model() |
| 873 | + |
| 874 | + stream = await provider.stream(model, [], system_prompt="system prompt here") |
| 875 | + _ = [e async for e in stream] |
| 876 | + result = await stream.result() |
| 877 | + |
| 878 | + usage = result.usage |
| 879 | + assert usage is not None |
| 880 | + assert usage.cache_read_tokens == 0 |
| 881 | + assert usage.cache_write_tokens > 0 |
| 882 | + assert usage.input_tokens > 0 |
| 883 | + assert usage.output_tokens > 0 |
| 884 | + # On first call, input_tokens should equal the total prompt tokens |
| 885 | + # because there's nothing in the cache |
| 886 | + assert usage.input_tokens == usage.cache_write_tokens |
| 887 | + |
| 888 | + async def test_cache_usage_second_call_prefix_match(self): |
| 889 | + """Second call with identical context: cache_read covers the full prompt, |
| 890 | + input_tokens is reduced, cache_write is minimal.""" |
| 891 | + provider = FauxProvider() |
| 892 | + provider.set_responses( |
| 893 | + [faux_assistant_message("first"), faux_assistant_message("second")] |
| 894 | + ) |
| 895 | + model = self._make_model() |
| 896 | + msgs = [UserMessage(content=[TextContent(text="hello")])] |
| 897 | + |
| 898 | + # First call |
| 899 | + s1 = await provider.stream(model, msgs, system_prompt="sys") |
| 900 | + _ = [e async for e in s1] |
| 901 | + await s1.result() |
| 902 | + |
| 903 | + # Second call with exact same context |
| 904 | + s2 = await provider.stream(model, msgs, system_prompt="sys") |
| 905 | + _ = [e async for e in s2] |
| 906 | + r2 = await s2.result() |
| 907 | + |
| 908 | + assert r2.usage is not None |
| 909 | + # Full prefix match: all prompt tokens come from cache |
| 910 | + assert r2.usage.cache_read_tokens > 0 |
| 911 | + assert r2.usage.cache_write_tokens == 0 |
| 912 | + assert r2.usage.input_tokens == 0 |
| 913 | + |
| 914 | + async def test_cache_usage_partial_prefix_change(self): |
| 915 | + """When messages grow, the prefix still matches and cache_read is partial.""" |
| 916 | + provider = FauxProvider() |
| 917 | + provider.set_responses( |
| 918 | + [faux_assistant_message("first"), faux_assistant_message("second")] |
| 919 | + ) |
| 920 | + model = self._make_model() |
| 921 | + |
| 922 | + # First call: one message |
| 923 | + msgs1 = [UserMessage(content=[TextContent(text="hello")])] |
| 924 | + s1 = await provider.stream(model, msgs1, system_prompt="sys") |
| 925 | + _ = [e async for e in s1] |
| 926 | + |
| 927 | + # Second call: same prefix + additional message |
| 928 | + msgs2 = [ |
| 929 | + UserMessage(content=[TextContent(text="hello")]), |
| 930 | + UserMessage(content=[TextContent(text="world")]), |
| 931 | + ] |
| 932 | + s2 = await provider.stream(model, msgs2, system_prompt="sys") |
| 933 | + _ = [e async for e in s2] |
| 934 | + r2 = await s2.result() |
| 935 | + |
| 936 | + assert r2.usage is not None |
| 937 | + # Should have partial cache read (the common prefix) |
| 938 | + assert r2.usage.cache_read_tokens > 0 |
| 939 | + # Should have cache write for the new part |
| 940 | + assert r2.usage.cache_write_tokens > 0 |
0 commit comments