"""Smoke tests for DeepSpeed accelerate config detection."""
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from opsd_utils.deepspeed_utils import (
    deepspeed_requires_single_student_forward,
    deepspeed_zero_stage,
    gradient_checkpointing_enable_kwargs,
    is_deepspeed_accelerate_config,
    should_colocate_teacher_with_student,
    should_disable_gradient_checkpointing,
    student_forward_chunk_size,
    uses_deepspeed_json_file,
)


def test_zero2_config_detected(monkeypatch):
    monkeypatch.setenv("ACCELERATE_CONFIG", "default_config_zero2.yaml")
    assert is_deepspeed_accelerate_config()
    assert uses_deepspeed_json_file()
    assert deepspeed_zero_stage() == 2
    assert should_colocate_teacher_with_student("auto")
    assert gradient_checkpointing_enable_kwargs() == {"use_reentrant": False}
    assert deepspeed_requires_single_student_forward()
    assert should_disable_gradient_checkpointing()
    assert student_forward_chunk_size(32, has_vision=True) == 32
    assert student_forward_chunk_size(32, has_vision=False) == 32


def test_ddp_config_not_deepspeed(monkeypatch):
    monkeypatch.setenv("ACCELERATE_CONFIG", "default_config.yaml")
    assert not is_deepspeed_accelerate_config()
    assert not should_colocate_teacher_with_student("auto")
    assert not deepspeed_requires_single_student_forward()
    assert student_forward_chunk_size(8, has_vision=True) == 1