pig.py/tests/unit/test_neopig_state.py

"""
Tests for neopig.state module.

Tests AppendOnlyStateLog and state file helpers.
"""

import pytest
import tempfile
import shutil
import os
import json
from pathlib import Path

from neopig.state import (
    AppendOnlyStateLog,
    get_state_log_path,
    get_state_file_path,
    rotate_state_file,
)


class TestAppendOnlyStateLog:
    """Test AppendOnlyStateLog class."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.log_path = Path(self.temp_dir) / "test.log"

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    def test_context_manager(self):
        """Test context manager opens and closes file."""
        with AppendOnlyStateLog(self.log_path) as log:
            assert log._file is not None
        assert log._file is None

    def test_page_record(self):
        """Test recording page visits."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.page("https://example.com/page1")
            log.page("https://example.com/page2")

        content = self.log_path.read_text()
        assert "P https://example.com/page1" in content
        assert "P https://example.com/page2" in content

    def test_media_record(self):
        """Test recording media downloads."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.media("abc123", "https://example.com/image.jpg")

        content = self.log_path.read_text()
        assert "M abc123 https://example.com/image.jpg" in content

    def test_screenshot_record(self):
        """Test recording screenshot captures."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.screenshot("https://example.com/page")

        content = self.log_path.read_text()
        assert "S https://example.com/page" in content

    def test_skip_domain_record(self):
        """Test recording domains to skip."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.skip_domain("blocked.com")

        content = self.log_path.read_text()
        assert "D blocked.com" in content

    def test_stats_record(self):
        """Test recording stats checkpoint."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.stats({"pages": 100, "media": 50})

        content = self.log_path.read_text()
        assert "X stats" in content
        assert '"pages": 100' in content

    def test_load_empty(self):
        """Test loading from non-existent file."""
        log = AppendOnlyStateLog(self.log_path)
        result = log.load()

        assert result['seen_pages'] == set()
        assert result['seen_media'] == {}
        assert result['seen_screenshots'] == set()
        assert result['skip_domains'] == set()
        assert result['stats'] == {}

    def test_load_populated(self):
        """Test loading from populated file."""
        with AppendOnlyStateLog(self.log_path) as log:
            log.page("https://example.com/page1")
            log.media("hash1", "https://example.com/img.jpg")
            log.screenshot("https://example.com/page1")
            log.skip_domain("blocked.com")
            log.stats({"count": 5})

        log = AppendOnlyStateLog(self.log_path)
        result = log.load()

        assert "https://example.com/page1" in result['seen_pages']
        assert result['seen_media']["https://example.com/img.jpg"] == "hash1"
        assert "https://example.com/page1" in result['seen_screenshots']
        assert "blocked.com" in result['skip_domains']
        assert result['stats'] == {"count": 5}


class TestGetStateLogPath:
    """Test get_state_log_path function."""

    def test_basic_domain(self):
        """Test path generation for basic domain."""
        path = get_state_log_path("example.com")
        assert str(path) == "data/example-com.log"

    def test_domain_with_protocol(self):
        """Test path generation for domain with protocol."""
        path = get_state_log_path("https://example.com")
        assert "example-com" in str(path)


class TestGetStateFilePath:
    """Test get_state_file_path function."""

    def test_basic_domain(self):
        """Test path generation for basic domain."""
        path = get_state_file_path("example.com")
        assert str(path) == "data/example-com.state"


class TestRotateStateFile:
    """Test rotate_state_file function."""

    def setup_method(self):
        self.temp_dir = tempfile.mkdtemp()
        self.state_path = Path(self.temp_dir) / "test.state"

    def teardown_method(self):
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    def test_no_rotation_if_not_exists(self):
        """Test no rotation if file doesn't exist."""
        result = rotate_state_file(self.state_path)
        assert result is None

    def test_rotation_creates_numbered_file(self):
        """Test rotation creates numbered backup."""
        self.state_path.write_text('{"key": "value"}')

        rotated = rotate_state_file(self.state_path)

        assert rotated is not None
        assert rotated.exists()
        assert rotated.name == "test.state.1"
        assert not self.state_path.exists()

    def test_multiple_rotations(self):
        """Test multiple rotations increment number."""
        self.state_path.write_text('{"key": "value1"}')
        rotate_state_file(self.state_path)

        self.state_path.write_text('{"key": "value2"}')
        rotated = rotate_state_file(self.state_path)

        assert rotated.name == "test.state.2"

    def test_preserve_keys(self):
        """Test preserving specific keys during rotation."""
        self.state_path.write_text('{"keep": "this", "drop": "that"}')

        rotate_state_file(self.state_path, preserve_keys=["keep"])

        assert self.state_path.exists()
        preserved = json.loads(self.state_path.read_text())
        assert preserved == {"keep": "this"}


if __name__ == '__main__':
    pytest.main([__file__, '-v'])