hud-python 0.4.7__py3-none-any.whl → 0.4.9__py3-none-any.whl

hud/cli/tests/test_push.py

@@ -0,0 +1,379 @@
+"""Tests for push.py - Push HUD environments to registry."""
+
+from __future__ import annotations
+
+import base64
+import json
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from unittest import mock
+
+import pytest
+import requests
+import typer
+import yaml
+
+from hud.cli.push import (
+    get_docker_image_labels,
+    get_docker_username,
+    push_command,
+    push_environment,
+)
+
+
+class TestGetDockerUsername:
+    """Test getting Docker username."""
+
+    def test_get_username_from_config(self, tmp_path):
+        """Test getting username from Docker config."""
+        # Create mock Docker config
+        docker_dir = tmp_path / ".docker"
+        docker_dir.mkdir()
+        
+        config_file = docker_dir / "config.json"
+        config = {
+            "auths": {
+                "https://index.docker.io/v1/": {
+                    "auth": base64.b64encode(b"testuser:testpass").decode()
+                }
+            }
+        }
+        config_file.write_text(json.dumps(config))
+        
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            username = get_docker_username()
+        
+        assert username == "testuser"
+
+    def test_get_username_no_config(self, tmp_path):
+        """Test when no Docker config exists."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            username = get_docker_username()
+        
+        assert username is None
+
+    def test_get_username_token_auth(self, tmp_path):
+        """Test skipping token-based auth."""
+        docker_dir = tmp_path / ".docker"
+        docker_dir.mkdir()
+        
+        config_file = docker_dir / "config.json"
+        config = {
+            "auths": {
+                "docker.io": {
+                    "auth": base64.b64encode(b"token:xyz").decode()
+                }
+            }
+        }
+        config_file.write_text(json.dumps(config))
+        
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            username = get_docker_username()
+        
+        assert username is None
+
+    @mock.patch("subprocess.run")
+    def test_get_username_credential_helper(self, mock_run, tmp_path):
+        """Test getting username from credential helper."""
+        docker_dir = tmp_path / ".docker"
+        docker_dir.mkdir()
+        
+        config_file = docker_dir / "config.json"
+        config = {"credsStore": "desktop"}
+        config_file.write_text(json.dumps(config))
+        
+        # Mock credential helper calls
+        mock_run.side_effect = [
+            mock.Mock(returncode=0, stdout='{"https://index.docker.io/v1/": "creds"}'),
+            mock.Mock(returncode=0, stdout='{"Username": "helperuser", "Secret": "pass"}')
+        ]
+        
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            username = get_docker_username()
+        
+        assert username == "helperuser"
+
+
+class TestGetDockerImageLabels:
+    """Test getting Docker image labels."""
+
+    @mock.patch("subprocess.run")
+    def test_get_labels_success(self, mock_run):
+        """Test successfully getting image labels."""
+        labels = {
+            "org.hud.manifest.head": "abc123",
+            "org.hud.version": "1.0.0"
+        }
+        mock_run.return_value = mock.Mock(
+            returncode=0,
+            stdout=json.dumps(labels),
+            stderr=""
+        )
+        
+        result = get_docker_image_labels("test:latest")
+        assert result == labels
+
+    @mock.patch("subprocess.run")
+    def test_get_labels_failure(self, mock_run):
+        """Test handling failure to get labels."""
+        mock_run.side_effect = Exception("Command failed")
+        
+        result = get_docker_image_labels("test:latest")
+        assert result == {}
+
+
+class TestPushEnvironment:
+    """Test the main push_environment function."""
+
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_no_lock_file(self, mock_design_class, tmp_path):
+        """Test pushing when no lock file exists."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        
+        with pytest.raises(typer.Exit) as exc_info:
+            push_environment(str(tmp_path))
+        
+        assert exc_info.value.exit_code == 1
+        mock_design.error.assert_called()
+
+    @mock.patch("hud.cli.push.HUDDesign")
+    @mock.patch("hud.cli.push.settings")
+    def test_push_no_api_key(self, mock_settings, mock_design_class, tmp_path):
+        """Test pushing without API key."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        mock_settings.api_key = None
+        
+        # Create lock file
+        lock_file = tmp_path / "hud.lock.yaml"
+        lock_file.write_text(yaml.dump({"image": "test:latest"}))
+        
+        with pytest.raises(typer.Exit) as exc_info:
+            push_environment(str(tmp_path))
+        
+        assert exc_info.value.exit_code == 1
+
+    @mock.patch("requests.post")
+    @mock.patch("subprocess.Popen")
+    @mock.patch("subprocess.run")
+    @mock.patch("hud.cli.push.get_docker_username")
+    @mock.patch("hud.cli.push.settings")
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_auto_detect_username(
+        self, mock_design_class, mock_settings, mock_get_username,
+        mock_run, mock_popen, mock_post, tmp_path
+    ):
+        """Test auto-detecting Docker username and pushing."""
+        # Setup mocks
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        mock_settings.api_key = "test-key"
+        mock_settings.hud_telemetry_url = "https://api.hud.test"
+        mock_get_username.return_value = "testuser"
+        
+        # Create lock file
+        lock_data = {
+            "image": "original/image:v1.0",
+            "build": {"version": "0.1.0"}
+        }
+        lock_file = tmp_path / "hud.lock.yaml"
+        lock_file.write_text(yaml.dump(lock_data))
+        
+        # Mock docker commands
+        def mock_run_impl(*args, **kwargs):
+            cmd = args[0]
+            if cmd[1] == "inspect":
+                if len(cmd) == 3:  # docker inspect <image>
+                    return mock.Mock(returncode=0, stdout="")
+                else:  # docker inspect --format ... <image>
+                    return mock.Mock(returncode=0, stdout="testuser/image:0.1.0@sha256:abc123")
+            elif cmd[1] == "tag":
+                return mock.Mock(returncode=0)
+            return mock.Mock(returncode=0)
+        
+        mock_run.side_effect = mock_run_impl
+        
+        # Mock docker push
+        mock_process = mock.Mock()
+        mock_process.stdout = ["Pushing image...", "Push complete"]
+        mock_process.wait.return_value = None
+        mock_process.returncode = 0
+        mock_popen.return_value = mock_process
+        
+        # Mock registry upload
+        mock_post.return_value = mock.Mock(status_code=201)
+        
+        # Run push
+        push_environment(str(tmp_path), yes=True)
+        
+        # Verify docker commands
+        assert mock_run.call_count >= 2
+        mock_popen.assert_called_once()
+        
+        # Verify registry upload
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        assert "testuser/image:0.1.0" in call_args[0][0]
+
+    @mock.patch("subprocess.run")
+    @mock.patch("hud.cli.push.settings")
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_explicit_image(
+        self, mock_design_class, mock_settings, mock_run, tmp_path
+    ):
+        """Test pushing with explicit image name."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        mock_settings.api_key = "test-key"
+        
+        # Create lock file
+        lock_data = {"image": "local:latest"}
+        lock_file = tmp_path / "hud.lock.yaml"
+        lock_file.write_text(yaml.dump(lock_data))
+        
+        # Mock docker inspect for non-existent local image
+        mock_run.side_effect = subprocess.CalledProcessError(1, "docker")
+        
+        with pytest.raises(typer.Exit):
+            push_environment(str(tmp_path), image="myrepo/myimage:v2")
+
+    @mock.patch("subprocess.Popen")
+    @mock.patch("subprocess.run")
+    @mock.patch("hud.cli.push.settings")
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_with_tag(
+        self, mock_design_class, mock_settings, mock_run, mock_popen, tmp_path
+    ):
+        """Test pushing with explicit tag."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        mock_settings.api_key = "test-key"
+        
+        # Create lock file
+        lock_data = {"image": "test:latest"}
+        lock_file = tmp_path / "hud.lock.yaml"
+        lock_file.write_text(yaml.dump(lock_data))
+        
+        # Mock docker commands
+        def mock_run_impl(*args, **kwargs):
+            cmd = args[0]
+            if cmd[1] == "inspect":
+                if len(cmd) == 3:  # docker inspect <image>
+                    return mock.Mock(returncode=0)
+                else:  # docker inspect --format ... <image>
+                    return mock.Mock(returncode=0, stdout="user/test:v2.0")
+            elif cmd[1] == "tag":
+                return mock.Mock(returncode=0)
+            return mock.Mock(returncode=0)
+        
+        mock_run.side_effect = mock_run_impl
+        
+        # Mock docker push
+        mock_process = mock.Mock()
+        mock_process.stdout = []
+        mock_process.wait.return_value = None
+        mock_process.returncode = 0
+        mock_popen.return_value = mock_process
+        
+        # Run push
+        push_environment(str(tmp_path), image="user/test", tag="v2.0", yes=True)
+        
+        # Verify tag was used
+        tag_call = [c for c in mock_run.call_args_list if c[0][0][1] == "tag"]
+        assert len(tag_call) > 0
+        assert "user/test:v2.0" in tag_call[0][0][0]
+
+    @mock.patch("subprocess.Popen")
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_docker_failure(self, mock_design_class, mock_popen):
+        """Test handling Docker push failure."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        
+        # Mock docker push failure
+        mock_process = mock.Mock()
+        mock_process.stdout = ["Error: access denied"]
+        mock_process.wait.return_value = None
+        mock_process.returncode = 1
+        mock_popen.return_value = mock_process
+        
+        with mock.patch("hud.cli.push.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            with mock.patch("subprocess.run"):
+                with pytest.raises(typer.Exit):
+                    push_environment(".", image="test:latest", yes=True)
+
+    @mock.patch("hud.cli.push.get_docker_image_labels")
+    @mock.patch("subprocess.run")
+    @mock.patch("hud.cli.push.settings")
+    @mock.patch("hud.cli.push.HUDDesign")
+    def test_push_with_labels(
+        self, mock_design_class, mock_settings, mock_run, mock_get_labels, tmp_path
+    ):
+        """Test pushing with image labels."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        mock_settings.api_key = "test-key"
+        
+        # Create lock file
+        lock_data = {"image": "test:latest"}
+        lock_file = tmp_path / "hud.lock.yaml"
+        lock_file.write_text(yaml.dump(lock_data))
+        
+        # Mock labels
+        mock_get_labels.return_value = {
+            "org.hud.manifest.head": "abc123def456",
+            "org.hud.version": "1.2.3"
+        }
+        
+        # Mock docker commands - first inspect succeeds to get to label check
+        # Provide explicit image to bypass username check
+        def mock_run_impl(*args, **kwargs):
+            cmd = args[0]
+            if cmd[1] == "inspect" and len(cmd) == 3:
+                # First inspect to check if image exists
+                return mock.Mock(returncode=0)
+            elif cmd[1] == "tag":
+                # Fail on tag to exit after labels are checked
+                raise subprocess.CalledProcessError(1, cmd)
+            return mock.Mock(returncode=0)
+        
+        mock_run.side_effect = mock_run_impl
+        
+        # Provide explicit image to ensure we reach label check
+        with pytest.raises(subprocess.CalledProcessError):
+            push_environment(str(tmp_path), image="test:v2", verbose=True)
+        
+        # Verify labels were checked
+        mock_get_labels.assert_called_once_with("test:latest")
+
+
+class TestPushCommand:
+    """Test the CLI command wrapper."""
+
+    def test_push_command_basic(self):
+        """Test basic push command."""
+        with mock.patch("hud.cli.push.push_environment") as mock_push:
+            push_command()
+            
+            mock_push.assert_called_once_with(
+                ".", None, None, False, False, False
+            )
+
+    def test_push_command_with_options(self):
+        """Test push command with all options."""
+        with mock.patch("hud.cli.push.push_environment") as mock_push:
+            push_command(
+                directory="./myenv",
+                image="myrepo/myimage",
+                tag="v1.0",
+                sign=True,
+                yes=True,
+                verbose=True
+            )
+            
+            mock_push.assert_called_once_with(
+                "./myenv", "myrepo/myimage", "v1.0", True, True, True
+            )

hud/cli/tests/test_registry.py

@@ -0,0 +1,264 @@
+"""Tests for registry.py - Local registry management for HUD environments."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest import mock
+
+import pytest
+import yaml
+
+from hud.cli.registry import (
+    extract_digest_from_image,
+    extract_name_and_tag,
+    get_registry_dir,
+    list_registry_entries,
+    load_from_registry,
+    save_to_registry,
+)
+
+
+class TestGetRegistryDir:
+    """Test getting registry directory."""
+
+    def test_get_registry_dir(self):
+        """Test default registry directory."""
+        with mock.patch("pathlib.Path.home") as mock_home:
+            mock_home.return_value = Path("/home/user")
+            
+            registry_dir = get_registry_dir()
+            assert registry_dir == Path("/home/user/.hud/envs")
+
+
+class TestExtractDigestFromImage:
+    """Test extracting digest from Docker image reference."""
+
+    def test_extract_from_full_digest(self):
+        """Test extracting from full digest reference."""
+        image = "myimage:tag@sha256:abc123def456789"
+        digest = extract_digest_from_image(image)
+        assert digest == "abc123def456"
+
+    def test_extract_from_digest_only(self):
+        """Test extracting from digest-only format."""
+        image = "sha256:deadbeef1234567890"
+        digest = extract_digest_from_image(image)
+        assert digest == "deadbeef1234"
+
+    def test_extract_from_tag(self):
+        """Test extracting from tagged image."""
+        image = "myimage:v1.2.3"
+        digest = extract_digest_from_image(image)
+        assert digest == "v1.2.3"
+
+    def test_extract_from_long_tag(self):
+        """Test extracting from long tag (truncated)."""
+        image = "myimage:superlongtagname123456789"
+        digest = extract_digest_from_image(image)
+        assert digest == "superlongtag"  # Max 12 chars
+
+    def test_extract_no_tag(self):
+        """Test extracting from image without tag."""
+        image = "myimage"
+        digest = extract_digest_from_image(image)
+        assert digest == "latest"
+
+    def test_extract_with_registry(self):
+        """Test extracting from image with registry."""
+        image = "docker.io/library/ubuntu:20.04"
+        digest = extract_digest_from_image(image)
+        assert digest == "20.04"
+
+    def test_extract_with_port(self):
+        """Test extracting from image with port."""
+        image = "localhost:5000/myimage"
+        digest = extract_digest_from_image(image)
+        assert digest == "latest"
+
+
+class TestExtractNameAndTag:
+    """Test extracting name and tag from Docker image reference."""
+
+    def test_extract_simple(self):
+        """Test extracting from simple image reference."""
+        name, tag = extract_name_and_tag("myorg/myapp:v1.0")
+        assert name == "myorg/myapp"
+        assert tag == "v1.0"
+
+    def test_extract_with_digest(self):
+        """Test extracting from reference with digest."""
+        name, tag = extract_name_and_tag("docker.io/hudpython/test:latest@sha256:abc123")
+        assert name == "hudpython/test"
+        assert tag == "latest"
+
+    def test_extract_no_tag(self):
+        """Test extracting from reference without tag."""
+        name, tag = extract_name_and_tag("myorg/myapp")
+        assert name == "myorg/myapp"
+        assert tag == "latest"
+
+    def test_extract_with_docker_registry(self):
+        """Test extracting from reference with docker.io prefix."""
+        name, tag = extract_name_and_tag("docker.io/library/ubuntu:20.04")
+        assert name == "library/ubuntu"
+        assert tag == "20.04"
+
+    def test_extract_with_other_registry(self):
+        """Test extracting from reference with custom registry."""
+        name, tag = extract_name_and_tag("gcr.io/myproject/myapp:v2")
+        assert name == "gcr.io/myproject/myapp"
+        assert tag == "v2"
+
+    def test_extract_single_name(self):
+        """Test extracting from single name without org."""
+        name, tag = extract_name_and_tag("ubuntu")
+        assert name == "ubuntu"
+        assert tag == "latest"
+
+
+class TestSaveToRegistry:
+    """Test saving to local registry."""
+
+    @mock.patch("hud.cli.registry.HUDDesign")
+    def test_save_success(self, mock_design_class, tmp_path):
+        """Test successful save to registry."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        
+        # Mock home directory
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            lock_data = {
+                "image": "test:latest@sha256:abc123",
+                "tools": ["tool1", "tool2"]
+            }
+            
+            result = save_to_registry(lock_data, "test:latest@sha256:abc123def456789")
+            
+            assert result is not None
+            assert result.exists()
+            assert result.name == "hud.lock.yaml"
+            
+            # Verify content
+            with open(result) as f:
+                saved_data = yaml.safe_load(f)
+            assert saved_data == lock_data
+            
+            # Verify directory structure
+            assert result.parent.name == "abc123def456"
+            
+            mock_design.success.assert_called_once()
+
+    @mock.patch("hud.cli.registry.HUDDesign")
+    def test_save_verbose(self, mock_design_class, tmp_path):
+        """Test save with verbose output."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            lock_data = {"image": "test:v1"}
+            
+            result = save_to_registry(lock_data, "test:v1", verbose=True)
+            
+            assert result is not None
+            # Should show verbose info
+            assert mock_design.info.call_count >= 1
+
+    @mock.patch("hud.cli.registry.HUDDesign")
+    def test_save_failure(self, mock_design_class):
+        """Test handling save failure."""
+        mock_design = mock.Mock()
+        mock_design_class.return_value = mock_design
+        
+        # Mock file operations to fail
+        with mock.patch("builtins.open", side_effect=IOError("Permission denied")):
+            with mock.patch("pathlib.Path.home", return_value=Path("/tmp")):
+                lock_data = {"image": "test:latest"}
+                
+                result = save_to_registry(lock_data, "test:latest", verbose=True)
+                
+                assert result is None
+                mock_design.warning.assert_called_once()
+
+
+class TestLoadFromRegistry:
+    """Test loading from local registry."""
+
+    def test_load_success(self, tmp_path):
+        """Test successful load from registry."""
+        # Create test registry structure
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            registry_dir = get_registry_dir()
+            digest_dir = registry_dir / "abc123"
+            digest_dir.mkdir(parents=True)
+            
+            lock_data = {"image": "test:latest", "version": "1.0"}
+            lock_file = digest_dir / "hud.lock.yaml"
+            lock_file.write_text(yaml.dump(lock_data))
+            
+            # Load it back
+            loaded = load_from_registry("abc123")
+            assert loaded == lock_data
+
+    def test_load_not_found(self, tmp_path):
+        """Test loading non-existent entry."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            loaded = load_from_registry("nonexistent")
+            assert loaded is None
+
+    def test_load_corrupted(self, tmp_path):
+        """Test loading corrupted lock file."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            registry_dir = get_registry_dir()
+            digest_dir = registry_dir / "bad"
+            digest_dir.mkdir(parents=True)
+            
+            lock_file = digest_dir / "hud.lock.yaml"
+            lock_file.write_text("invalid: yaml: content:")
+            
+            loaded = load_from_registry("bad")
+            assert loaded is None
+
+
+class TestListRegistryEntries:
+    """Test listing registry entries."""
+
+    def test_list_empty(self, tmp_path):
+        """Test listing empty registry."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            entries = list_registry_entries()
+            assert entries == []
+
+    def test_list_entries(self, tmp_path):
+        """Test listing multiple entries."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path):
+            registry_dir = get_registry_dir()
+            
+            # Create several entries
+            for digest in ["abc123", "def456", "ghi789"]:
+                digest_dir = registry_dir / digest
+                digest_dir.mkdir(parents=True)
+                lock_file = digest_dir / "hud.lock.yaml"
+                lock_file.write_text(f"image: test:{digest}")
+            
+            # Create a directory without lock file (should be ignored)
+            (registry_dir / "nolockfile").mkdir(parents=True)
+            
+            # Create a file in registry dir (should be ignored)
+            (registry_dir / "README.txt").write_text("info")
+            
+            entries = list_registry_entries()
+            
+            assert len(entries) == 3
+            digests = [entry[0] for entry in entries]
+            assert set(digests) == {"abc123", "def456", "ghi789"}
+            
+            # Verify all paths are lock files
+            for _, lock_path in entries:
+                assert lock_path.name == "hud.lock.yaml"
+                assert lock_path.exists()
+
+    def test_list_no_registry_dir(self, tmp_path):
+        """Test listing when registry directory doesn't exist."""
+        with mock.patch("pathlib.Path.home", return_value=tmp_path / "nonexistent"):
+            entries = list_registry_entries()
+            assert entries == []

hud/clients/base.py

@@ -101,6 +101,7 @@ class BaseHUDClient(AgentMCPClient):
         self._mcp_config = mcp_config
         self._strict_validation = strict_validation
         self._auto_trace = auto_trace
+        self._auto_trace_cm: Any | None = None  # Store auto-created trace context manager
 
         self._initialized = False
         self._telemetry_data = {}  # Initialize telemetry data
@@ -124,7 +125,7 @@ class BaseHUDClient(AgentMCPClient):
                 "Either pass it to the constructor or call initialize with a configuration"
             )
 
-        setup_hud_telemetry(self._mcp_config, auto_trace=self._auto_trace)
+        self._auto_trace_cm = setup_hud_telemetry(self._mcp_config, auto_trace=self._auto_trace)
 
         logger.debug("Initializing MCP client...")
 
@@ -163,6 +164,17 @@ class BaseHUDClient(AgentMCPClient):
 
     async def shutdown(self) -> None:
         """Disconnect from the MCP server."""
+        # Clean up auto-created trace if any
+        if self._auto_trace_cm:
+            try:
+                self._auto_trace_cm.__exit__(None, None, None)
+                logger.info("Closed auto-created trace")
+            except Exception as e:
+                logger.warning("Failed to close auto-created trace: %s", e)
+            finally:
+                self._auto_trace_cm = None
+        
+        # Disconnect from server
         if self._initialized:
             await self._disconnect()
             self._initialized = False