tianxuan/tests/test_data_loader.py

"""Tests for analysis.data_loader — BOM detection, schema validation, load/config."""

import csv
import os
import tempfile
from pathlib import Path

import pytest
import polars as pl

from analysis.data_loader import (
    _file_count,
    _merge_schema_entries,
    _resolve_encoding,
    detect_bom,
    load_config,
    load_csv_directory,
    validate_schemas,
)


# ═══════════════════════════════════════════════════════════════════════
#  BOM detection
# ═══════════════════════════════════════════════════════════════════════

class TestDetectBOM:
    """detect_bom() should return the correct encoding string for each BOM."""

    def test_utf8_bom(self):
        """UTF-8 BOM (EF BB BF) → 'utf-8-sig'."""
        path = _write_bytes(b'\xef\xbb\xbfcol1,col2\n1,2\n')
        try:
            assert detect_bom(path) == 'utf-8-sig'
        finally:
            os.unlink(path)

    def test_utf16le_bom(self):
        """UTF-16 LE BOM (FF FE) → 'utf-16le'."""
        path = _write_bytes(b'\xff\xfec\x00o\x00l\x001\x00,\x00c\x00o\x00l\x002\x00\n\x001\x00,\x002\x00\n')
        try:
            assert detect_bom(path) == 'utf-16le'
        finally:
            os.unlink(path)

    def test_utf16be_bom(self):
        """UTF-16 BE BOM (FE FF) → 'utf-16be'."""
        path = _write_bytes(b'\xfe\xff\x00c\x00o\x00l\x001\x00,\x00c\x00o\x00l\x002\x00\n\x001\x00,\x002\x00\n')
        try:
            assert detect_bom(path) == 'utf-16be'
        finally:
            os.unlink(path)

    def test_no_bom(self):
        """Plain ASCII/UTF-8 without BOM → 'utf-8'."""
        path = _write_bytes(b'col1,col2\n1,2\n')
        try:
            assert detect_bom(path) == 'utf-8'
        finally:
            os.unlink(path)

    def test_empty_file(self):
        """Empty file → 'utf-8'."""
        path = _write_bytes(b'')
        try:
            assert detect_bom(path) == 'utf-8'
        finally:
            os.unlink(path)


# ═══════════════════════════════════════════════════════════════════════
#  Schema validation
# ═══════════════════════════════════════════════════════════════════════

class TestValidateSchemas:
    """validate_schemas() must raise ValueError when column sets differ."""

    def test_matching_schemas(self):
        """Identical schemas → no exception."""
        schemas = [
            {':ips': 'Utf8', '8ack': 'Int64'},
            {':ips': 'Utf8', '8ack': 'Int64'},
        ]
        validate_schemas(schemas)  # should not raise

    def test_empty_list(self):
        """Empty list → no exception."""
        validate_schemas([])

    def test_schema_mismatch_missing_column(self):
        """One file missing a column → ValueError."""
        schemas = [
            {':ips': 'Utf8', '8ack': 'Int64'},
            {':ips': 'Utf8'},  # missing '8ack'
        ]
        with pytest.raises(ValueError, match=r'missing columns'):
            validate_schemas(schemas)

    def test_schema_mismatch_extra_column(self):
        """One file has an extra column → ValueError."""
        schemas = [
            {':ips': 'Utf8'},
            {':ips': 'Utf8', 'extra_col': 'Int64'},
        ]
        with pytest.raises(ValueError, match=r'extra columns'):
            validate_schemas(schemas)

    def test_schema_mismatch_both_sides(self):
        """Both missing and extra columns → ValueError with both messages."""
        schemas = [
            {'a': 'Int64', 'b': 'Utf8'},
            {'a': 'Int64', 'c': 'Float64'},
        ]
        with pytest.raises(ValueError):
            validate_schemas(schemas)


class TestMergeSchemaEntries:
    """_merge_schema_entries() should merge without raising."""

    def test_simple_merge(self):
        merged, conflicts = _merge_schema_entries([
            {'a': 'Int64', 'b': 'Utf8'},
            {'a': 'Int64', 'c': 'Float64'},
        ])
        assert set(merged.keys()) == {'a', 'b', 'c'}
        assert merged['a'] == 'Int64'

    def test_type_conflict(self):
        """Different dtypes for same column → last wins, conflict recorded."""
        merged, conflicts = _merge_schema_entries([
            {'val': 'Int64'},
            {'val': 'Float64'},
        ])
        assert merged['val'] == 'Float64'
        assert 'val' in conflicts


# ═══════════════════════════════════════════════════════════════════════
#  CSV loading
# ═══════════════════════════════════════════════════════════════════════

class TestLoadCSVDirectory:
    """End-to-end load_csv_directory() tests with temporary CSV files."""

    def test_basic_load(self):
        """Single CSV → LazyFrame + correct schema + row/file counts."""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / 'test.csv'
            _write_csv(path, [
                [':ips', ':ipd', '8ack'],
                ['10.0.0.1', '10.0.0.2', '100'],
                ['10.0.0.1', '10.0.0.3', '200'],
                ['10.0.0.2', '10.0.0.4', '150'],
            ])

            lf, schema, row_count, file_count, memory_mb = load_csv_directory(
                str(path), encoding='utf-8'
            )

            assert isinstance(lf, pl.LazyFrame)
            assert ':ips' in schema
            assert ':ipd' in schema
            assert '8ack' in schema
            assert file_count == 1
            assert row_count > 0
            assert memory_mb >= 0

            # Verify data round-trip
            df = lf.collect()
            assert len(df) == 3
            assert set(df[':ips'].to_list()) == {'10.0.0.1', '10.0.0.2'}

    def test_multiple_files(self):
        """Multiple CSVs with same schema → merged LazyFrame."""
        with tempfile.TemporaryDirectory() as tmpdir:
            p1 = Path(tmpdir) / 'part1.csv'
            p2 = Path(tmpdir) / 'part2.csv'
            _write_csv(p1, [['x'], ['1'], ['2']])
            _write_csv(p2, [['x'], ['3'], ['4']])

            glob_pattern = str(Path(tmpdir) / 'part*.csv')
            lf, schema, row_count, file_count, memory_mb = load_csv_directory(
                glob_pattern, encoding='utf-8'
            )

            assert file_count == 2
            df = lf.collect()
            assert len(df) == 4
            assert sorted(df['x'].to_list()) == [1, 2, 3, 4]

    def test_no_files(self):
        """No matching files → FileNotFoundError."""
        with pytest.raises(FileNotFoundError, match=r'No files match'):
            load_csv_directory('__nonexistent_glarg_*.csv')

    def test_schema_mismatch_across_files_strict(self):
        """Files with different column sets + schema_strict=True → ValueError."""
        with tempfile.TemporaryDirectory() as tmpdir:
            _write_csv(Path(tmpdir) / 'a.csv', [['x', 'y'], ['1', '2']])
            _write_csv(Path(tmpdir) / 'b.csv', [['x'], ['3']])

            with pytest.raises(ValueError, match=r'Schema mismatch'):
                load_csv_directory(str(Path(tmpdir) / '*.csv'), schema_strict=True)

    def test_schema_mismatch_across_files_lenient(self):
        """Files with different column sets + schema_strict=False → merge with nulls."""
        with tempfile.TemporaryDirectory() as tmpdir:
            _write_csv(Path(tmpdir) / 'a.csv', [['x', 'y'], ['1', '2']])
            _write_csv(Path(tmpdir) / 'b.csv', [['x'], ['3']])

            lf, schema, *_ = load_csv_directory(str(Path(tmpdir) / '*.csv'))
            df = lf.collect()
            assert len(df) == 2
            assert 'x' in df.columns
            assert 'y' in df.columns
            # Row from b.csv should have null for 'y'
            assert df['y'].is_null().sum() == 1

    def test_schema_override(self):
        """schema_override forces a specific column dtype."""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / 'test.csv'
            _write_csv(path, [['val'], ['123'], ['456']])

            lf, schema, *_ = load_csv_directory(
                str(path), encoding='utf-8',
                schema_override={'val': 'Float64'},
            )

            df = lf.collect()
            assert df['val'].dtype == pl.Float64

    def test_custom_delimiter(self):
        """Tab-delimited file should parse correctly."""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / 'test.tsv'
            with open(path, 'w', newline='') as f:
                f.write('a\tb\n1\t2\n3\t4\n')

            lf, *_ = load_csv_directory(str(path), delimiter='\t')
            df = lf.collect()
            assert len(df) == 2
            assert df.columns == ['a', 'b']

    def test_glob_with_subdirectory(self):
        """Recursive glob should find files in subdirectories."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sub = Path(tmpdir) / 'sub'
            sub.mkdir()
            _write_csv(sub / 'nested.csv', [['x'], ['42']])

            lf, schema, row_count, file_count, memory_mb = load_csv_directory(
                str(Path(tmpdir) / '**' / '*.csv')
            )
            assert file_count == 1
            assert row_count > 0

    def test_latlon_with_plus_and_blank(self):
        """CSV with lat/lon columns containing '+', '' loads without error
        and the column is Float64 after cleaning."""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / 'latlon_test.csv'
            with open(path, 'w', newline='') as f:
                f.write('src_ip,latitude,longitude\n')
                f.write('10.0.0.1,+,120.0\n')
                f.write('10.0.0.2,35.5,\n')
                f.write('10.0.0.3,-,\n')
                f.write('10.0.0.4,40.0,121.5\n')

            lf, schema, *_ = load_csv_directory(str(path))
            df = lf.collect()

            # lat/lon columns should be Float64 in schema
            assert 'latitude' in schema
            assert 'longitude' in schema
            assert 'latitude' in df.columns
            assert 'longitude' in df.columns

            # Check dtypes are Float64
            assert df['latitude'].dtype == pl.Float64
            assert df['longitude'].dtype == pl.Float64

            # Standalone '+' and blank should become null
            assert df['latitude'].is_null().sum() == 2  # rows 0,2
            assert df['longitude'].is_null().sum() == 2  # rows 1,2

            # Valid values survive
            assert df['latitude'].drop_nulls().to_list() == [35.5, 40.0]
            assert df['longitude'].drop_nulls().to_list() == [120.0, 121.5]


# ═══════════════════════════════════════════════════════════════════════
#  Config loading
# ═══════════════════════════════════════════════════════════════════════

class TestLoadConfig:
    """load_config() YAML loading and error handling."""

    def test_load_valid_yaml(self):
        path = _write_yaml('key: value\nnested:\n  inner: 42\n')
        try:
            config = load_config(path)
            assert config['key'] == 'value'
            assert config['nested']['inner'] == 42
        finally:
            os.unlink(path)

    def test_load_empty_yaml(self):
        path = _write_yaml('')
        try:
            config = load_config(path)
            assert config == {}
        finally:
            os.unlink(path)

    def test_config_path_none(self):
        assert load_config(None) == {}

    def test_config_not_found(self):
        with pytest.raises(FileNotFoundError):
            load_config('/__nonexistent__/config.yaml')


# ═══════════════════════════════════════════════════════════════════════
#  Utility functions
# ═══════════════════════════════════════════════════════════════════════

class TestResolveEncoding:
    def test_bom_overrides_default(self):
        assert _resolve_encoding('utf-8', 'utf-16le') == 'utf-16le'

    def test_no_bom_uses_default(self):
        # Polars >= 1.0 uses 'utf8' (not 'utf-8')
        assert _resolve_encoding('utf-8', 'utf-8') == 'utf8'
        assert _resolve_encoding('latin-1', 'utf-8') == 'latin-1'


class TestFileCount:
    def test_count(self):
        assert _file_count(['a', 'b', 'c']) == 3
        assert _file_count([]) == 0


# ═══════════════════════════════════════════════════════════════════════
#  Helpers
# ═══════════════════════════════════════════════════════════════════════

def _write_csv(path: Path, rows: list[list[str]]) -> None:
    """Write rows as a CSV file."""
    with open(path, 'w', newline='') as f:
        writer = csv.writer(f)
        for row in rows:
            writer.writerow(row)


def _write_bytes(data: bytes) -> str:
    """Write raw bytes to a temp file; return the path string."""
    with tempfile.NamedTemporaryFile(suffix='.csv', delete=False) as f:
        f.write(data)
        return f.name


def _write_yaml(content: str) -> str:
    """Write YAML string to a temp file; return the path string."""
    with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w', delete=False) as f:
        f.write(content)
        return f.name