| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | from unittest.mock import MagicMock |
| 16 | from colab_cli.runtime import ColabRuntime |
| 17 | |
| 18 | |
| 19 | def test_runtime_execute_code_streaming(): |
| 20 | runtime = ColabRuntime("http://url", "token") |
| 21 | mock_client = MagicMock() |
| 22 | |
| 23 | # Mock the return value of execute_interactive (the raw reply) |
| 24 | mock_client.execute_interactive.return_value = {"content": {"status": "ok"}} |
| 25 | |
| 26 | # Inject our mock client |
| 27 | runtime._kernel_client = mock_client |
| 28 | |
| 29 | streamed_outputs = [] |
| 30 | |
| 31 | def output_hook(o): |
| 32 | streamed_outputs.append(o) |
| 33 | |
| 34 | code = "print(1); print(2)" |
| 35 | |
| 36 | # We need to simulate the execution where execute_interactive is called |
| 37 | # and it calls our wrapped_output_hook. |
| 38 | # Since wrapped_output_hook is defined inside execute_code, we have to |
| 39 | # intercept the call to execute_interactive to get a reference to it. |
| 40 | |
| 41 | def side_effect(code, output_hook=None, **kwargs): |
| 42 | # Simulate messages arriving |
| 43 | msg1 = { |
| 44 | "header": {"msg_type": "stream"}, |
| 45 | "content": {"name": "stdout", "text": "1\n"}, |
| 46 | } |
| 47 | msg2 = { |
| 48 | "header": {"msg_type": "stream"}, |
| 49 | "content": {"name": "stdout", "text": "2\n"}, |
| 50 | } |
| 51 | if output_hook: |
| 52 | output_hook(msg1) |
| 53 | output_hook(msg2) |
| 54 | return {"content": {"status": "ok"}} |
| 55 | |
| 56 | mock_client.execute_interactive.side_effect = side_effect |
| 57 | |
| 58 | outputs = runtime.execute_code(code, output_hook=output_hook) |
| 59 | |
| 60 | assert len(outputs) == 2 |
| 61 | assert outputs[0]["text"] == "1\n" |
| 62 | assert outputs[1]["text"] == "2\n" |
| 63 | |
| 64 | # Verify streaming hook was called |
| 65 | assert len(streamed_outputs) == 2 |
| 66 | assert streamed_outputs[0]["text"] == "1\n" |
| 67 | assert streamed_outputs[1]["text"] == "2\n" |
| 68 | |
| 69 | |
| 70 | def test_runtime_execute_code_streaming_error_synthesis(): |
| 71 | runtime = ColabRuntime("http://url", "token") |
| 72 | mock_client = MagicMock() |
| 73 | |
| 74 | # Simulate an error reply but NO error output message |
| 75 | mock_client.execute_interactive.return_value = { |
| 76 | "content": { |
| 77 | "status": "error", |
| 78 | "ename": "RuntimeError", |
| 79 | "evalue": "something went wrong", |
| 80 | "traceback": ["tb line 1"], |
| 81 | } |
| 82 | } |
| 83 | runtime._kernel_client = mock_client |
| 84 | |
| 85 | streamed_outputs = [] |
| 86 | outputs = runtime.execute_code( |
| 87 | "fail()", output_hook=lambda o: streamed_outputs.append(o) |
| 88 | ) |
| 89 | |
| 90 | # Final outputs should include synthesized error |
| 91 | assert len(outputs) == 1 |
| 92 | assert outputs[0]["output_type"] == "error" |
| 93 | assert outputs[0]["ename"] == "RuntimeError" |
| 94 | |
| 95 | # Note: synthesized error is added AFTER execute_interactive returns, |
| 96 | # so it won't be in streamed_outputs unless we specifically add logic for it. |
| 97 | # Currently it's only in the returned list. |
| 98 | assert len(streamed_outputs) == 0 |