dstack/src/tests/_internal/cli/utils/test_common.py at cb8ec5ce0a6818cfa7453d80c040bafa402da651 · dstackai/dstack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Generator
from unittest.mock import patch

import pytest

from dstack._internal.cli.utils.common import _get_cli_log_file


@pytest.fixture
def mock_dstack_dir(tmp_path: Path) -> Generator[Path, None, None]:
    with patch("dstack._internal.cli.utils.common.get_dstack_dir") as mock:
        mock.return_value = tmp_path
        yield tmp_path


class TestGetCliLogFile:
    def test_no_existing_dir(self, mock_dstack_dir: Path):
        log_dir = mock_dstack_dir / "logs" / "cli"
        expected_log_file = log_dir / "latest.log"
        assert not log_dir.exists()

        result = _get_cli_log_file()

        assert log_dir.exists()
        assert result == expected_log_file

    def test_no_rotation_needed_for_today_file(self, mock_dstack_dir: Path):
        log_dir = mock_dstack_dir / "logs" / "cli"
        log_dir.mkdir(parents=True, exist_ok=True)
        latest_log = log_dir / "latest.log"
        latest_log.touch()

        result = _get_cli_log_file()

        assert result == latest_log
        assert latest_log.exists(), "latest.log should not have been renamed"

    @patch("dstack._internal.cli.utils.common.datetime")
    def test_simple_rotation(self, mock_datetime, mock_dstack_dir: Path):
        # Mock "now" to be a specific date
        now = datetime(2023, 10, 27, 10, 0, 0, tzinfo=timezone.utc)
        mock_datetime.now.return_value = now
        # Ensure fromtimestamp still works correctly for the System Under Test
        mock_datetime.fromtimestamp.side_effect = lambda ts, tz: datetime.fromtimestamp(ts, tz)

        log_dir = mock_dstack_dir / "logs" / "cli"
        log_dir.mkdir(parents=True, exist_ok=True)
        latest_log = log_dir / "latest.log"
        latest_log.touch()

        # Set the modification time to yesterday
        yesterday = now - timedelta(days=1)
        mtime = yesterday.timestamp()
        os.utime(latest_log, (mtime, mtime))

        # The expected rotated file name is based on the modification time (yesterday)
        date_str = yesterday.strftime("%Y-%m-%d")
        expected_rotated_log = log_dir / f"{date_str}.log"

        result = _get_cli_log_file()

        assert result == log_dir / "latest.log"
        assert not latest_log.exists(), "The original latest.log should have been renamed"
        assert expected_rotated_log.exists(), "The log file should have been rotated"

    @patch("dstack._internal.cli.utils.common.datetime")
    def test_rotation_with_conflict(self, mock_datetime, mock_dstack_dir: Path):
        now = datetime(2023, 10, 27, 10, 0, 0, tzinfo=timezone.utc)
        yesterday = now - timedelta(days=1)
        mock_datetime.now.return_value = now
        mock_datetime.fromtimestamp.side_effect = lambda ts, tz: datetime.fromtimestamp(ts, tz)

        log_dir = mock_dstack_dir / "logs" / "cli"
        log_dir.mkdir(parents=True, exist_ok=True)

        # Create the old 'latest.log' and set its modification time to yesterday
        latest_log = log_dir / "latest.log"
        latest_log.touch()
        mtime = yesterday.timestamp()
        os.utime(latest_log, (mtime, mtime))

        # Create conflicting files that already exist from a previous rotation
        date_str = yesterday.strftime("%Y-%m-%d")
        conflicting_log_1 = log_dir / f"{date_str}.log"
        conflicting_log_1.touch()
        conflicting_log_2 = log_dir / f"{date_str}-1.log"
        conflicting_log_2.touch()

        # We expect the file to be rotated to the next available counter
        expected_rotated_log = log_dir / f"{date_str}-2.log"

        result = _get_cli_log_file()

        assert result == log_dir / "latest.log"
        assert not latest_log.exists(), "The original latest.log should have been renamed"
        assert conflicting_log_1.exists(), "Existing rotated log should be untouched"
        assert conflicting_log_2.exists(), "Existing rotated log with counter should be untouched"
        assert expected_rotated_log.exists(), (
            "The log should have rotated to the next available counter"
        )