| 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 | import multiprocessing |
| 16 | import os |
| 17 | import pytest |
| 18 | import tempfile |
| 19 | import threading |
| 20 | from datetime import datetime |
| 21 | |
| 22 | import filelock |
| 23 | |
| 24 | from colab_cli.state import StateStore, SessionState, SettingsStore, Settings |
| 25 | |
| 26 | |
| 27 | @pytest.fixture |
| 28 | def temp_config(): |
| 29 | with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f: |
| 30 | path = f.name |
| 31 | yield path |
| 32 | if os.path.exists(path): |
| 33 | os.remove(path) |
| 34 | |
| 35 | |
| 36 | def test_state_store_add_get(temp_config): |
| 37 | store = StateStore(temp_config) |
| 38 | state = SessionState( |
| 39 | name="test-session", |
| 40 | token="token123", |
| 41 | url="http://localhost", |
| 42 | endpoint="endpoint456", |
| 43 | variant="TPU", |
| 44 | accelerator="V5E1", |
| 45 | ) |
| 46 | store.add(state) |
| 47 | |
| 48 | # Reload store |
| 49 | new_store = StateStore(temp_config) |
| 50 | loaded = new_store.get("test-session") |
| 51 | assert loaded is not None |
| 52 | assert loaded.name == "test-session" |
| 53 | assert loaded.token == "token123" |
| 54 | assert loaded.variant == "TPU" |
| 55 | |
| 56 | |
| 57 | def test_state_store_remove(temp_config): |
| 58 | store = StateStore(temp_config) |
| 59 | state = SessionState(name="to-be-removed", token="tok", url="url", endpoint="end") |
| 60 | store.add(state) |
| 61 | assert store.get("to-be-removed") is not None |
| 62 | |
| 63 | store.remove("to-be-removed") |
| 64 | assert store.get("to-be-removed") is None |
| 65 | |
| 66 | # Reload check |
| 67 | new_store = StateStore(temp_config) |
| 68 | assert new_store.get("to-be-removed") is None |
| 69 | |
| 70 | |
| 71 | def test_state_store_list(temp_config): |
| 72 | store = StateStore(temp_config) |
| 73 | s1 = SessionState(name="s1", token="t1", url="u1", endpoint="e1") |
| 74 | s2 = SessionState(name="s2", token="t2", url="u2", endpoint="e2") |
| 75 | store.add(s1) |
| 76 | store.add(s2) |
| 77 | |
| 78 | sessions = store.list() |
| 79 | assert len(sessions) == 2 |
| 80 | assert "s1" in sessions |
| 81 | assert "s2" in sessions |
| 82 | |
| 83 | |
| 84 | def test_state_store_invalid_json(temp_config): |
| 85 | with open(temp_config, "w") as f: |
| 86 | f.write("invalid json") |
| 87 | |
| 88 | store = StateStore(temp_config) |
| 89 | assert store.list() == {} |
| 90 | |
| 91 | |
| 92 | def test_state_store_concurrency(temp_config): |
| 93 | def add_sessions(start, count, path): |
| 94 | store = StateStore(path) |
| 95 | for i in range(start, start + count): |
| 96 | s = SessionState(name=f"s{i}", token="t", url="u", endpoint="e") |
| 97 | store.add(s) |
| 98 | |
| 99 | t1 = threading.Thread(target=add_sessions, args=(0, 50, temp_config)) |
| 100 | t2 = threading.Thread(target=add_sessions, args=(50, 50, temp_config)) |
| 101 | |
| 102 | t1.start() |
| 103 | t2.start() |
| 104 | t1.join() |
| 105 | t2.join() |
| 106 | |
| 107 | new_store = StateStore(temp_config) |
| 108 | # This might pass or fail depending on luck without locking |
| 109 | # But usually it fails with 100 iterations. |
| 110 | assert len(new_store.list()) == 100 |
| 111 | |
| 112 | |
| 113 | def test_settings_store_defaults(temp_config): |
| 114 | store = SettingsStore(temp_config) |
| 115 | settings = store.load() |
| 116 | assert settings.enable_update_check is True |
| 117 | assert settings.last_check is None |
| 118 | |
| 119 | |
| 120 | def test_settings_store_save_load(temp_config): |
| 121 | store = SettingsStore(temp_config) |
| 122 | settings = Settings(enable_update_check=False, last_check=datetime(2026, 1, 1)) |
| 123 | store.save(settings) |
| 124 | |
| 125 | loaded = SettingsStore(temp_config).load() |
| 126 | assert loaded.enable_update_check is False |
| 127 | assert loaded.last_check == datetime(2026, 1, 1) |
| 128 | |
| 129 | |
| 130 | # --- Cross-platform locking (filelock ReadWriteLock) behavior --------------- |
| 131 | |
| 132 | |
| 133 | def test_lock_path_is_derived_from_path(temp_config): |
| 134 | """The lock sidecar file lives next to the data file with a .lock suffix.""" |
| 135 | store = StateStore(temp_config) |
| 136 | assert store.lock_path == temp_config + ".lock" |
| 137 | |
| 138 | |
| 139 | def test_store_uses_readwrite_lock_on_sidecar(temp_config): |
| 140 | """The store holds a ReadWriteLock bound to the sidecar path.""" |
| 141 | store = StateStore(temp_config) |
| 142 | assert isinstance(store._rwlock, filelock.ReadWriteLock) |
| 143 | assert store._rwlock.lock_file == temp_config + ".lock" |
| 144 | |
| 145 | |
| 146 | def test_separate_store_instances_can_write_from_threads(temp_config): |
| 147 | """Two StateStore instances writing from different threads must not crash. |
| 148 | |
| 149 | Regression guard for filelock's is_singleton=True default: it merges two |
| 150 | same-path ReadWriteLock objects into one reentrant lock whose reentrancy |
| 151 | guard raises RuntimeError when two threads contend for the write lock. We |
| 152 | pass is_singleton=False so they serialize via the underlying file lock. |
| 153 | """ |
| 154 | errors = [] |
| 155 | |
| 156 | def writer(start): |
| 157 | store = StateStore(temp_config) |
| 158 | try: |
| 159 | for i in range(start, start + 25): |
| 160 | store.add(SessionState(name=f"n{i}", token="t", url="u", endpoint="e")) |
| 161 | except Exception as exc: # pragma: no cover - failure path |
| 162 | errors.append(repr(exc)) |
| 163 | |
| 164 | threads = [threading.Thread(target=writer, args=(s,)) for s in (0, 100)] |
| 165 | for t in threads: |
| 166 | t.start() |
| 167 | for t in threads: |
| 168 | t.join() |
| 169 | |
| 170 | assert errors == [] |
| 171 | assert len(StateStore(temp_config).list()) == 50 |
| 172 | |
| 173 | |
| 174 | def test_write_acquires_write_lock(temp_config, mocker): |
| 175 | """Writes go through the ReadWriteLock's exclusive write_lock().""" |
| 176 | store = StateStore(temp_config) |
| 177 | spy = mocker.spy(store._rwlock, "write_lock") |
| 178 | store.add(SessionState(name="s", token="t", url="u", endpoint="e")) |
| 179 | assert spy.call_count >= 1 |
| 180 | assert store.get("s") is not None |
| 181 | |
| 182 | |
| 183 | def test_read_acquires_read_lock(temp_config, mocker): |
| 184 | """Reads go through the ReadWriteLock's shared read_lock().""" |
| 185 | store = StateStore(temp_config) |
| 186 | store.add(SessionState(name="s", token="t", url="u", endpoint="e")) |
| 187 | |
| 188 | spy = mocker.spy(store._rwlock, "read_lock") |
| 189 | assert store.get("s") is not None |
| 190 | assert spy.call_count >= 1 |
| 191 | |
| 192 | |
| 193 | def test_uses_filelock_not_fcntl(temp_config): |
| 194 | """The store must lock via the platform-independent filelock library. |
| 195 | |
| 196 | Guards against a regression back to the POSIX-only fcntl.flock approach |
| 197 | that broke Windows users. |
| 198 | """ |
| 199 | store = StateStore(temp_config) |
| 200 | assert store._rwlock.__class__.__module__.startswith("filelock") |
| 201 | |
| 202 | |
| 203 | def test_no_fcntl_import_in_state_module(): |
| 204 | """The state module must not depend on the POSIX-only fcntl module.""" |
| 205 | import colab_cli.state as state_module |
| 206 | |
| 207 | assert not hasattr(state_module, "fcntl") |
| 208 | |
| 209 | |
| 210 | def _mp_hold_write_lock(lock_path, hold_for, acquired_evt, release_evt): |
| 211 | rw = filelock.ReadWriteLock(lock_path) |
| 212 | with rw.write_lock(): |
| 213 | acquired_evt.set() |
| 214 | release_evt.wait(timeout=hold_for) |
| 215 | |
| 216 | |
| 217 | def test_write_lock_blocks_across_processes(temp_config): |
| 218 | """A write lock held by another process must block the store's write section. |
| 219 | |
| 220 | The cross-process guarantee is the whole point of the fcntl->filelock |
| 221 | switch, so we hold the write lock from a separate process and confirm the |
| 222 | in-process writer cannot proceed until it's released. |
| 223 | """ |
| 224 | store = StateStore(temp_config) |
| 225 | # Pre-create the data file so add() doesn't race on first creation. |
| 226 | store.add(SessionState(name="seed", token="t", url="u", endpoint="e")) |
| 227 | |
| 228 | ctx = multiprocessing.get_context("spawn") |
| 229 | acquired = ctx.Event() |
| 230 | release = ctx.Event() |
| 231 | holder = ctx.Process( |
| 232 | target=_mp_hold_write_lock, args=(store.lock_path, 10, acquired, release) |
| 233 | ) |
| 234 | holder.start() |
| 235 | try: |
| 236 | assert acquired.wait(timeout=5), "holder process never acquired the lock" |
| 237 | |
| 238 | finished = threading.Event() |
| 239 | |
| 240 | def writer(): |
| 241 | store.add(SessionState(name="blocked", token="t", url="u", endpoint="e")) |
| 242 | finished.set() |
| 243 | |
| 244 | t = threading.Thread(target=writer) |
| 245 | t.start() |
| 246 | |
| 247 | # While the external process holds the write lock, the writer can't finish. |
| 248 | assert not finished.wait(timeout=0.75) |
| 249 | |
| 250 | release.set() |
| 251 | assert finished.wait(timeout=5) |
| 252 | t.join() |
| 253 | assert store.get("blocked") is not None |
| 254 | finally: |
| 255 | release.set() |
| 256 | holder.join(timeout=5) |
| 257 | |
| 258 | |
| 259 | def _mp_hold_read_lock(lock_path, acquired_evt, release_evt): |
| 260 | rw = filelock.ReadWriteLock(lock_path) |
| 261 | with rw.read_lock(): |
| 262 | acquired_evt.set() |
| 263 | release_evt.wait(timeout=10) |
| 264 | |
| 265 | |
| 266 | def test_readers_are_concurrent_across_processes(temp_config): |
| 267 | """Shared read locks must allow concurrent readers (the LOCK_SH semantics). |
| 268 | |
| 269 | This is the capability ReadWriteLock buys us over a plain exclusive |
| 270 | FileLock: while one process holds a read lock, this process can still read. |
| 271 | """ |
| 272 | store = StateStore(temp_config) |
| 273 | store.add(SessionState(name="s", token="t", url="u", endpoint="e")) |
| 274 | |
| 275 | ctx = multiprocessing.get_context("spawn") |
| 276 | acquired = ctx.Event() |
| 277 | release = ctx.Event() |
| 278 | holder = ctx.Process( |
| 279 | target=_mp_hold_read_lock, args=(store.lock_path, acquired, release) |
| 280 | ) |
| 281 | holder.start() |
| 282 | try: |
| 283 | assert acquired.wait(timeout=5), "holder process never acquired the read lock" |
| 284 | |
| 285 | done = threading.Event() |
| 286 | result = {} |
| 287 | |
| 288 | def reader(): |
| 289 | result["value"] = store.get("s") |
| 290 | done.set() |
| 291 | |
| 292 | t = threading.Thread(target=reader) |
| 293 | t.start() |
| 294 | |
| 295 | # The read must complete even though another process holds a read lock. |
| 296 | assert done.wait(timeout=3), "concurrent read was blocked by another reader" |
| 297 | t.join() |
| 298 | assert result["value"] is not None |
| 299 | finally: |
| 300 | release.set() |
| 301 | holder.join(timeout=5) |
| 302 | |
| 303 | |
| 304 | def _mp_add_sessions(path, start, count): |
| 305 | store = StateStore(path) |
| 306 | for i in range(start, start + count): |
| 307 | store.add(SessionState(name=f"p{i}", token="t", url="u", endpoint="e")) |
| 308 | |
| 309 | |
| 310 | def test_state_store_multiprocess_concurrency(temp_config): |
| 311 | """filelock must serialize writes across separate processes (not just threads). |
| 312 | |
| 313 | fcntl.flock is per-open-file-description and advisory; this test exercises |
| 314 | the cross-process guarantee that motivated the switch. |
| 315 | """ |
| 316 | ctx = multiprocessing.get_context("spawn") |
| 317 | p1 = ctx.Process(target=_mp_add_sessions, args=(temp_config, 0, 40)) |
| 318 | p2 = ctx.Process(target=_mp_add_sessions, args=(temp_config, 40, 40)) |
| 319 | |
| 320 | p1.start() |
| 321 | p2.start() |
| 322 | p1.join() |
| 323 | p2.join() |
| 324 | |
| 325 | assert p1.exitcode == 0 |
| 326 | assert p2.exitcode == 0 |
| 327 | assert len(StateStore(temp_config).list()) == 80 |