@samitouri / QOSami-colab-cli
Code Commits Issues Pulls Sessions CI/CD
main
main
py 165 lines 5.52 KB
Raw
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, patch
16
17 import jupyter_kernel_client
18
19 from colab_cli.runtime import ColabRuntime
20
21
22 def test_colab_runtime_kernel_client():
23 target_attr = "ColabKernelClient" if hasattr(jupyter_kernel_client, "ColabKernelClient") else "KernelClient"
24 token_param_name = "proxy_token" if hasattr(jupyter_kernel_client, "ColabKernelClient") else "token"
25
26 with patch.object(jupyter_kernel_client, target_attr) as mock_kc_cls:
27 mock_kc = mock_kc_cls.return_value
28 runtime = ColabRuntime("http://url", "token123")
29
30 assert runtime._kernel_client is None
31
32 kc = runtime.kernel_client
33
34 expected_kwargs = {
35 "server_url": "http://url",
36 token_param_name: "token123",
37 "kernel_id": None,
38 "client_kwargs": {
39 "subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT,
40 "extra_params": {"colab-runtime-proxy-token": "token123"},
41 },
42 "headers": {
43 "X-Colab-Client-Agent": "colab-cli",
44 "X-Colab-Runtime-Proxy-Token": "token123",
45 },
46 }
47 mock_kc_cls.assert_called_once_with(**expected_kwargs)
48 mock_kc.start.assert_called_once()
49 assert kc == mock_kc
50
51
52 def test_colab_runtime_execute_code():
53 runtime = ColabRuntime("http://url", "token123")
54 mock_kc = MagicMock()
55 runtime._kernel_client = mock_kc
56
57 # Test empty reply
58 mock_kc.execute.return_value = {}
59 assert runtime.execute_code("print(1)") == []
60
61 # Test normal reply
62 mock_kc.execute.return_value = {"outputs": [{"text": "1\n"}]}
63 assert runtime.execute_code("print(1)") == [{"text": "1\n"}]
64
65 # Test error status without error output
66 mock_kc.execute.return_value = {
67 "status": "error",
68 "ename": "ValueError",
69 "evalue": "bad",
70 "outputs": [{"text": "partial"}],
71 }
72 outputs = runtime.execute_code("raise ValueError")
73 assert len(outputs) == 2
74 assert outputs[0] == {"text": "partial"}
75 assert outputs[1] == {
76 "output_type": "error",
77 "ename": "ValueError",
78 "evalue": "bad",
79 "traceback": [],
80 }
81
82
83 def test_colab_runtime_execute_code_default_no_timeout():
84 """By default, execute_code should NOT pass a timeout (relies on jupyter
85 kernel client default), preserving existing behavior for fast / streaming
86 workloads."""
87 runtime = ColabRuntime("http://url", "token123")
88 mock_kc = MagicMock()
89 runtime._kernel_client = mock_kc
90
91 mock_kc.execute.return_value = {"outputs": []}
92 runtime.execute_code("print(1)")
93
94 _, kwargs = mock_kc.execute.call_args
95 assert "timeout" not in kwargs
96
97
98 def test_colab_runtime_execute_code_with_timeout():
99 """When a timeout is supplied, it must be forwarded to kernel_client.execute."""
100 runtime = ColabRuntime("http://url", "token123")
101 mock_kc = MagicMock()
102 runtime._kernel_client = mock_kc
103
104 mock_kc.execute.return_value = {"outputs": []}
105 runtime.execute_code("print(1)", timeout=600)
106
107 _, kwargs = mock_kc.execute.call_args
108 assert kwargs.get("timeout") == 600
109
110
111 def test_colab_runtime_execute_interactive_with_timeout():
112 """timeout must also be plumbed through the execute_interactive branch
113 (used when an output_hook is supplied)."""
114 runtime = ColabRuntime("http://url", "token123")
115 mock_kc = MagicMock()
116 runtime._kernel_client = mock_kc
117
118 mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}}
119 runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600)
120
121 _, kwargs = mock_kc.execute_interactive.call_args
122 assert kwargs.get("timeout") == 600
123
124
125 def test_colab_runtime_stop():
126 runtime = ColabRuntime("http://url", "token123")
127 mock_kc = MagicMock()
128 runtime._kernel_client = mock_kc
129
130 runtime.stop()
131 mock_kc._manager.client.stop_channels.assert_called_once()
132
133
134 def test_colab_runtime_stop_exception(caplog):
135 runtime = ColabRuntime("http://url", "token123")
136 mock_kc = MagicMock()
137 mock_kc._manager.client.stop_channels.side_effect = Exception("Stop failed")
138 runtime._kernel_client = mock_kc
139
140 runtime.stop() # Should not raise
141 assert "Error stopping kernel client" in caplog.text
142
143
144 def test_colab_runtime_stdin_logging():
145 mock_history = MagicMock()
146 runtime = ColabRuntime(
147 "http://url", "token", session_name="test-s", history=mock_history
148 )
149 mock_kc = MagicMock()
150 runtime._kernel_client = mock_kc
151
152 mock_kc.execute.side_effect = lambda code, allow_stdin=False, stdin_hook=None: {
153 "outputs": [{"text": stdin_hook("Enter something: ")}]
154 }
155
156 with patch("colab_cli.runtime.input", return_value="user input"):
157 outputs = runtime.execute_code("code", allow_stdin=True)
158
159 assert outputs == [{"text": "user input"}]
160 mock_history.log_event.assert_any_call(
161 "test-s", "stdin_request", {"prompt": "Enter something: "}
162 )
163 mock_history.log_event.assert_any_call(
164 "test-s", "input_reply", {"value": "user input"}
165 )