Added tests for FHIVibes input class

import configparser

import tempfile

import pytest

from pathlib import Path

from gims.inputs import inputs_cls

from gims.inputs.FHIVibes import FHIVibes

from gims.prepare_input import write_input_files

SPECIES_DIR = 'gims/static/data/species_defaults/'

# Diamond Si unit cell

STRUCTURE = {

    "cell": [[5.43, 0, 0], [0, 5.43, 0], [0, 0, 5.43]],

    "positions": [

        {"position": [0.0, 0.0, 0.0], "species": "Si", "initMoment": 0, "constraint": False, "charge": 0},

        {"position": [1.3575, 1.3575, 1.3575], "species": "Si", "initMoment": 0, "constraint": False, "charge": 0},

    ],

}

BASE_FORM = {

    "spectroscopy": {},

    "basic": {"species": ["Si"], "xc": "pbe", "basisSettings": "light"},

    "k_grid": {"k_grid": [4, 4, 4]},

}

def _make_data(extra_sections=None):

    return {

        "code": "FHIaims",

        "workflow": "Spectroscopy",

        "structure": STRUCTURE,

        "form": {**BASE_FORM, **(extra_sections or {})},

    }

def _read_ini(path):

    config = configparser.ConfigParser()

    config.optionxform = str  # preserve key case

    config.read(path)

    return config

class TestFHIVibes:

    def test_routing(self):

        """Spectroscopy workflow is routed to FHIVibes, not FHIaims."""

        assert inputs_cls(_make_data()) is FHIVibes

    def test_phonopy_sections(self):

        """Phonopy workflow produces [phonopy] section and correct calculator settings."""

        data = _make_data({'phonopy': {'supercell_matrix': [2, 2, 2], 'displacement': 0.01}})

        with tempfile.TemporaryDirectory(dir='.') as tmpdir:

            tdir, tar_file, _ = write_input_files('FHIaims', data, tmpdir, SPECIES_DIR)

            assert tar_file == 'input_files.tar'

            base = Path(tdir) / 'tar' / 'input_files'

            assert (base / 'geometry.in').exists()

            cfg = _read_ini(base / 'vibes.in')

        assert cfg['files']['geometry'] == 'geometry.in'

        assert cfg['calculator']['name'] == 'aims'

        assert 'phonopy' in cfg

        assert cfg['calculator.parameters']['xc'] == 'pbe'

        assert cfg['calculator.basissets']['default'] == 'light'

    def test_md_sections(self):

        """MD workflow produces [md] and [md.kwargs] sections."""

        data = _make_data({'md_vibes': {'MD_params': [300, 0.01], 'dt': 1.0}})

        with tempfile.TemporaryDirectory(dir='.') as tmpdir:

            tdir, tar_file, _ = write_input_files('FHIaims', data, tmpdir, SPECIES_DIR)

            assert tar_file == 'input_files.tar'

            cfg = _read_ini(Path(tdir) / 'tar' / 'input_files' / 'vibes.in')

        assert 'md' in cfg

        assert 'md.kwargs' in cfg

        assert cfg['md.kwargs']['temperature'] == '300'

    def test_polarization_output_lines(self):

        """Polarization section produces one 'output: polarization N ...' line per direction."""

        pol = [30, 16, 1, 16, 30, 1, 16, 16, 30]

        data = _make_data({'polarization': {'output_polarization': pol}})

        with tempfile.TemporaryDirectory(dir='.') as tmpdir:

            tdir, tar_file, _ = write_input_files('FHIaims', data, tmpdir, SPECIES_DIR)

            assert tar_file == 'input_files.tar'

            text = (Path(tdir) / 'tar' / 'input_files' / 'vibes.in').read_text()

        assert 'output: polarization 1 30 16 1' in text

        assert 'output: polarization 2 16 30 1' in text

        assert 'output: polarization 3 16 16 30' in text

    def test_missing_basis_raises(self):

        """Missing basisSettings raises ValueError before writing any file."""

        data = _make_data()

        del data['form']['basic']['basisSettings']

        with tempfile.TemporaryDirectory(dir='.') as tmpdir:

            with pytest.raises(ValueError, match='basis set'):

                write_input_files('FHIaims', data, tmpdir, SPECIES_DIR)