1from contextlib import nullcontext as does_not_raise
2from datetime import datetime
3from itertools import cycle
4from typing import Callable, Generator, Iterable, Optional, Tuple, Union
5
6import yaml
7from freezegun import freeze_time
8from lava.utils.log_section import (
9    DEFAULT_GITLAB_SECTION_TIMEOUTS,
10    FALLBACK_GITLAB_SECTION_TIMEOUT,
11    LogSectionType,
12)
13
14
15def section_timeout(section_type: LogSectionType) -> int:
16    return int(
17        DEFAULT_GITLAB_SECTION_TIMEOUTS.get(
18            section_type, FALLBACK_GITLAB_SECTION_TIMEOUT
19        ).total_seconds()
20    )
21
22
23def create_lava_yaml_msg(
24    dt: Callable = datetime.now, msg="test", lvl="target"
25) -> dict[str, str]:
26    return {"dt": str(dt()), "msg": msg, "lvl": lvl}
27
28
29def generate_testsuite_result(
30    name="test-mesa-ci", result="pass", metadata_extra=None, extra=None
31):
32    if metadata_extra is None:
33        metadata_extra = {}
34    if extra is None:
35        extra = {}
36    return {"metadata": {"result": result, **metadata_extra}, "name": name}
37
38
39def jobs_logs_response(
40    finished=False, msg=None, lvl="target", result=None
41) -> Tuple[bool, str]:
42    timed_msg = {"dt": str(datetime.now()), "msg": "New message", "lvl": lvl}
43    if result:
44        timed_msg["lvl"] = "target"
45        timed_msg["msg"] = f"hwci: mesa: {result}"
46
47    logs = [timed_msg] if msg is None else msg
48
49    return finished, yaml.safe_dump(logs)
50
51
52def section_aware_message_generator(
53    messages: dict[LogSectionType, Iterable[int]], result: Optional[str] = None
54) -> Iterable[tuple[dict, Iterable[int]]]:
55    default = [1]
56
57    result_message_section = LogSectionType.TEST_CASE
58
59    for section_type in LogSectionType:
60        delay = messages.get(section_type, default)
61        yield mock_lava_signal(section_type), delay
62        if result and section_type == result_message_section:
63            # To consider the job finished, the result `echo` should be produced
64            # in the correct section
65            yield create_lava_yaml_msg(msg=f"hwci: mesa: {result}"), delay
66
67
68def message_generator():
69    for section_type in LogSectionType:
70        yield mock_lava_signal(section_type)
71
72
73def level_generator():
74    # Tests all known levels by default
75    yield from cycle(("results", "feedback", "warning", "error", "debug", "target"))
76
77
78def generate_n_logs(
79    n=1,
80    tick_fn: Union[Generator, Iterable[int], int] = 1,
81    level_fn=level_generator,
82    result="pass",
83):
84    """Simulate a log partitionated in n components"""
85    level_gen = level_fn()
86
87    if isinstance(tick_fn, Generator):
88        tick_gen = tick_fn
89    elif isinstance(tick_fn, Iterable):
90        tick_gen = cycle(tick_fn)
91    else:
92        tick_gen = cycle((tick_fn,))
93
94    with freeze_time(datetime.now()) as time_travel:
95        tick_sec: int = next(tick_gen)
96        while True:
97            # Simulate a scenario where the target job is waiting for being started
98            for _ in range(n - 1):
99                level: str = next(level_gen)
100
101                time_travel.tick(tick_sec)
102                yield jobs_logs_response(finished=False, msg=[], lvl=level)
103
104            time_travel.tick(tick_sec)
105            yield jobs_logs_response(finished=True, result=result)
106
107
108def to_iterable(tick_fn):
109    if isinstance(tick_fn, Generator):
110        return tick_fn
111    elif isinstance(tick_fn, Iterable):
112        return cycle(tick_fn)
113    else:
114        return cycle((tick_fn,))
115
116
117def mock_logs(messages=None, result=None):
118    if messages is None:
119        messages = {}
120    with freeze_time(datetime.now()) as time_travel:
121        # Simulate a complete run given by message_fn
122        for msg, tick_list in section_aware_message_generator(messages, result):
123            for tick_sec in tick_list:
124                yield jobs_logs_response(finished=False, msg=[msg])
125                time_travel.tick(tick_sec)
126
127
128def mock_lava_signal(type: LogSectionType) -> dict[str, str]:
129    return {
130        LogSectionType.TEST_CASE: create_lava_yaml_msg(
131            msg="<STARTTC> case", lvl="debug"
132        ),
133        LogSectionType.TEST_SUITE: create_lava_yaml_msg(
134            msg="<STARTRUN> suite", lvl="debug"
135        ),
136        LogSectionType.LAVA_POST_PROCESSING: create_lava_yaml_msg(
137            msg="<LAVA_SIGNAL_ENDTC case>", lvl="target"
138        ),
139    }.get(type, create_lava_yaml_msg())
140