- 
                Notifications
    
You must be signed in to change notification settings  - Fork 2.2k
 
Fix/transcript finish #3324
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Fix/transcript finish #3324
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | 
|---|---|---|
| 
          
            
          
           | 
    @@ -109,3 +109,53 @@ async def test_close(gemini_connection, mock_gemini_session): | |
| await gemini_connection.close() | ||
| 
     | 
||
| mock_gemini_session.close.assert_called_once() | ||
| 
     | 
||
| 
     | 
||
| @pytest.mark.asyncio | ||
| @pytest.mark.parametrize('tx_direction', ['input', 'output']) | ||
| async def test_receive_transcript_finished( | ||
| gemini_connection, mock_gemini_session, tx_direction | ||
| ): | ||
| """Test receive_transcript_finished for input and output transcription.""" | ||
| 
     | 
||
| finished_tx = types.Transcription(finished=True) | ||
| 
     | 
||
| class Msg: | ||
| 
     | 
||
| def __init__(self): | ||
| self.server_content = mock.Mock() | ||
| sc = self.server_content | ||
| sc.model_turn = None | ||
| if tx_direction == 'input': | ||
| sc.input_transcription = finished_tx | ||
| sc.output_transcription = None | ||
| else: | ||
| sc.input_transcription = None | ||
| sc.output_transcription = finished_tx | ||
| sc.interrupted = False | ||
| sc.turn_complete = False | ||
| self.tool_call = None | ||
| self.session_resumption_update = None | ||
| 
     | 
||
| async def gen(): | ||
| yield Msg() | ||
| 
     | 
||
| mock_gemini_session.receive = mock.Mock(return_value=gen()) | ||
| 
     | 
||
| responses = [] | ||
| async for r in gemini_connection.receive(): | ||
| responses.append(r) | ||
| 
     | 
||
| if tx_direction == 'input': | ||
| tx_resps = [r for r in responses if r.input_transcription] | ||
| else: | ||
| tx_resps = [r for r in responses if r.output_transcription] | ||
| 
     | 
||
| if tx_direction == 'input': | ||
| assert tx_resps, 'Excpected input transcription response' | ||
| assert tx_resps[0].input_transcription.finished is True | ||
| assert not tx_resps[0].input_transcription.text | ||
| else: | ||
| assert tx_resps, 'Expected output transcription response' | ||
| assert tx_resps[0].output_transcription.finished is True | ||
| assert not tx_resps[0].output_transcription.text | ||
| 
         
      Comment on lines
    
      +149
     to 
      +161
    
   
  There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The assertion logic is duplicated for the 'input' and 'output' cases. You can refactor this to be more DRY (Don't Repeat Yourself) by using a variable for the attribute name and consolidating the checks. This makes the test easier to read and maintain. This change also corrects a typo in the assertion message ('Excpected' -> 'Expected').   attr_name = f'{tx_direction}_transcription'
  tx_resps = [r for r in responses if getattr(r, attr_name)]
  assert tx_resps, f'Expected {tx_direction} transcription response'
  transcription = getattr(tx_resps[0], attr_name)
  assert transcription.finished is True
  assert not transcription.text | 
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For improved readability and conciseness, you can create the mock message object directly by setting attributes on a
mock.Mock()instance, instead of defining a localMsgclass. This makes the test setup more straightforward and less verbose.