Adds tests for run decorator, remote runner, and CLI prerequisites

Tests cover:
- run() decorator behavior, env var capture (exact, wildcard, mixed)
- Remote runner GCS helpers, execution flow, and error handling
- CLI prerequisite checks for gcloud, pulumi, kubectl, docker, and auth

Author: JyotinderSingh

keras_remote/cli/test_prerequisites_check.py

@@ -0,0 +1,67 @@
+"""Tests for keras_remote.cli.prerequisites_check — tool availability checks."""
+
+import click
+import pytest
+
+from keras_remote.cli.prerequisites_check import (
+  check_docker,
+  check_gcloud,
+  check_gcloud_auth,
+  check_kubectl,
+  check_pulumi,
+)
+
+
+@pytest.mark.parametrize(
+  ("check_fn", "error_match"),
+  [
+    (check_gcloud, "gcloud CLI not found"),
+    (check_pulumi, "Pulumi CLI not found"),
+    (check_kubectl, "kubectl not found"),
+    (check_docker, "Docker not found"),
+  ],
+  ids=["gcloud", "pulumi", "kubectl", "docker"],
+)
+class TestToolChecks:
+  def test_present(self, mocker, check_fn, error_match):
+    mocker.patch("shutil.which", return_value="/usr/bin/tool")
+    check_fn()
+
+  def test_missing(self, mocker, check_fn, error_match):
+    mocker.patch("shutil.which", return_value=None)
+    with pytest.raises(click.ClickException, match=error_match):
+      check_fn()
+
+
+class TestCheckGcloudAuth:
+  def test_token_success(self, mocker):
+    """When print-access-token succeeds, no login is triggered."""
+    mock_run = mocker.patch(
+      "keras_remote.cli.prerequisites_check.subprocess.run",
+    )
+    mock_run.return_value.returncode = 0
+
+    check_gcloud_auth()
+
+    # Only called once (the token check), not a second time for login
+    assert mock_run.call_count == 1
+
+  def test_token_failure_triggers_login(self, mocker):
+    """When print-access-token fails, gcloud auth login is run."""
+    mock_run = mocker.patch(
+      "keras_remote.cli.prerequisites_check.subprocess.run",
+    )
+    # First call = token check (fails), second call = login (succeeds)
+    token_result = mocker.MagicMock()
+    token_result.returncode = 1
+    mock_run.return_value = token_result
+
+    mocker.patch("keras_remote.cli.prerequisites_check.warning")
+    mocker.patch("click.echo")
+
+    check_gcloud_auth()
+
+    assert mock_run.call_count == 2
+    # Second call should be the login command
+    login_call = mock_run.call_args_list[1]
+    assert "login" in login_call[0][0]

keras_remote/core/test_core.py

@@ -0,0 +1,134 @@
+"""Tests for keras_remote.core.core — run decorator and env var capture."""
+
+from unittest.mock import MagicMock
+
+from keras_remote.core.core import run
+
+
+class TestRunDecorator:
+  def test_decorator_preserves_wrapped_function(self):
+    @run(accelerator="cpu")
+    def my_func():
+      """My docstring."""
+
+    assert callable(my_func)
+    assert my_func.__name__ == "my_func"
+    assert my_func.__doc__ == "My docstring."
+
+
+class TestEnvVarCapture:
+  def test_exact_match(self, monkeypatch, mocker):
+    monkeypatch.setenv("MY_VAR", "my_val")
+    mock_exec = mocker.patch("keras_remote.core.core._execute_on_gke")
+
+    @run(accelerator="cpu", capture_env_vars=["MY_VAR"])
+    def func():
+      pass
+
+    func()
+    call_args = mock_exec.call_args
+    env_vars = call_args[0][-1]  # last positional arg
+    assert env_vars == {"MY_VAR": "my_val"}
+
+  def test_wildcard_pattern(self, monkeypatch, mocker):
+    monkeypatch.setenv("PREFIX_A", "1")
+    monkeypatch.setenv("PREFIX_B", "2")
+    monkeypatch.setenv("OTHER", "3")
+    mock_exec = mocker.patch("keras_remote.core.core._execute_on_gke")
+
+    @run(accelerator="cpu", capture_env_vars=["PREFIX_*"])
+    def func():
+      pass
+
+    func()
+    env_vars = mock_exec.call_args[0][-1]
+    assert "PREFIX_A" in env_vars
+    assert "PREFIX_B" in env_vars
+    assert "OTHER" not in env_vars
+
+  def test_missing_var_skipped(self, monkeypatch, mocker):
+    monkeypatch.delenv("NONEXISTENT", raising=False)
+    mock_exec = mocker.patch("keras_remote.core.core._execute_on_gke")
+
+    @run(accelerator="cpu", capture_env_vars=["NONEXISTENT"])
+    def func():
+      pass
+
+    func()
+    env_vars = mock_exec.call_args[0][-1]
+    assert env_vars == {}
+
+  def test_none_capture(self, mocker):
+    mock_exec = mocker.patch("keras_remote.core.core._execute_on_gke")
+
+    @run(accelerator="cpu", capture_env_vars=None)
+    def func():
+      pass
+
+    func()
+    env_vars = mock_exec.call_args[0][-1]
+    assert env_vars == {}
+
+  def test_mixed_exact_and_wildcard(self, monkeypatch, mocker):
+    monkeypatch.setenv("EXACT_VAR", "exact")
+    monkeypatch.setenv("WILD_A", "a")
+    monkeypatch.setenv("WILD_B", "b")
+    mock_exec = mocker.patch("keras_remote.core.core._execute_on_gke")
+
+    @run(
+      accelerator="cpu",
+      capture_env_vars=["EXACT_VAR", "WILD_*"],
+    )
+    def func():
+      pass
+
+    func()
+    env_vars = mock_exec.call_args[0][-1]
+    assert env_vars == {"EXACT_VAR": "exact", "WILD_A": "a", "WILD_B": "b"}
+
+
+class TestExecuteOnGkeDefaults:
+  def test_cluster_from_env(self, monkeypatch, mocker):
+    """When cluster=None, falls back to KERAS_REMOTE_GKE_CLUSTER env var."""
+    monkeypatch.setenv("KERAS_REMOTE_GKE_CLUSTER", "env-cluster")
+    monkeypatch.setenv("KERAS_REMOTE_PROJECT", "proj")
+    mock_exec = mocker.patch(
+      "keras_remote.core.core.execute_remote",
+      return_value=42,
+    )
+    mocker.patch(
+      "keras_remote.core.core.JobContext.from_params",
+      return_value=MagicMock(),
+    )
+
+    @run(accelerator="cpu", cluster=None)
+    def func():
+      pass
+
+    func()
+
+    call_args = mock_exec.call_args
+    backend = call_args[0][1]
+    assert backend.cluster == "env-cluster"
+
+  def test_namespace_from_env(self, monkeypatch, mocker):
+    """When namespace=None, falls back to KERAS_REMOTE_GKE_NAMESPACE env var."""
+    monkeypatch.setenv("KERAS_REMOTE_GKE_NAMESPACE", "custom-ns")
+    monkeypatch.setenv("KERAS_REMOTE_PROJECT", "proj")
+    mock_exec = mocker.patch(
+      "keras_remote.core.core.execute_remote",
+      return_value=42,
+    )
+    mocker.patch(
+      "keras_remote.core.core.JobContext.from_params",
+      return_value=MagicMock(),
+    )
+
+    @run(accelerator="cpu", namespace=None)
+    def func():
+      pass
+
+    func()
+
+    backend = mock_exec.call_args[0][1]
+    assert backend.namespace == "custom-ns"