Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/py-polars/tests/unit/meta/test_cpu_check.py
8416 views
1
from __future__ import annotations

2


3
from typing import TYPE_CHECKING

4
from unittest.mock import Mock

5


6
import pytest

7


8
from polars import _cpu_check

9
from polars._cpu_check import check_cpu_flags

10


11
if TYPE_CHECKING:

12
    from tests.conftest import PlMonkeyPatch

13


14
TEST_FEATURE_FLAGS = "+sse3,+ssse3"

15


16


17
def test_check_cpu_flags(

18
    plmonkeypatch: PlMonkeyPatch, recwarn: pytest.WarningsRecorder

19
) -> None:

20
    cpu_flags = {"sse3": True, "ssse3": True}

21
    mock_read_cpu_flags = Mock(return_value=cpu_flags)

22
    plmonkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)

23


24
    check_cpu_flags(TEST_FEATURE_FLAGS)

25


26
    assert len(recwarn) == 0

27


28


29
def test_check_cpu_flags_missing_features(plmonkeypatch: PlMonkeyPatch) -> None:

30
    cpu_flags = {"sse3": True, "ssse3": False}

31
    mock_read_cpu_flags = Mock(return_value=cpu_flags)

32
    plmonkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)

33


34
    with pytest.warns(RuntimeWarning, match="Missing required CPU features") as w:

35
        check_cpu_flags(TEST_FEATURE_FLAGS)

36


37
    assert "ssse3" in str(w[0].message)

38


39


40
def test_check_cpu_flags_unknown_flag(

41
    plmonkeypatch: PlMonkeyPatch,

42
) -> None:

43
    real_cpu_flags = {"sse3": True, "ssse3": False}

44
    mock_read_cpu_flags = Mock(return_value=real_cpu_flags)

45
    unknown_feature_flags = "+sse3,+ssse3,+HelloWorld!"

46
    plmonkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)

47
    with pytest.raises(RuntimeError, match="unknown feature flag: 'HelloWorld!'"):

48
        check_cpu_flags(unknown_feature_flags)

49


50


51
def test_check_cpu_flags_skipped_no_flags(plmonkeypatch: PlMonkeyPatch) -> None:

52
    mock_read_cpu_flags = Mock()

53
    plmonkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)

54


55
    check_cpu_flags("")

56


57
    assert mock_read_cpu_flags.call_count == 0

58


59


60
def test_check_cpu_flags_skipped_env_var(plmonkeypatch: PlMonkeyPatch) -> None:

61
    plmonkeypatch.setenv("POLARS_SKIP_CPU_CHECK", "1")

62


63
    mock_read_cpu_flags = Mock()

64
    plmonkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)

65


66
    check_cpu_flags(TEST_FEATURE_FLAGS)

67


68
    assert mock_read_cpu_flags.call_count == 0

69


70