1
0
mirror of https://github.com/gryf/coach.git synced 2025-12-18 19:50:17 +01:00
Files
coach/rl_coach/tests/graph_managers/test_graph_manager.py

48 lines
1.9 KiB
Python

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
import pytest
from rl_coach.graph_managers.graph_manager import GraphManager, ScheduleParameters
from rl_coach.base_parameters import VisualizationParameters
from rl_coach.core_types import RunPhase
@pytest.mark.unit_test
def test_phase_context():
graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
assert graph_manager.phase == RunPhase.UNDEFINED
with graph_manager.phase_context(RunPhase.TRAIN):
assert graph_manager.phase == RunPhase.TRAIN
assert graph_manager.phase == RunPhase.UNDEFINED
@pytest.mark.unit_test
def test_phase_context_nested():
graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
assert graph_manager.phase == RunPhase.UNDEFINED
with graph_manager.phase_context(RunPhase.TRAIN):
assert graph_manager.phase == RunPhase.TRAIN
with graph_manager.phase_context(RunPhase.TEST):
assert graph_manager.phase == RunPhase.TEST
assert graph_manager.phase == RunPhase.TRAIN
assert graph_manager.phase == RunPhase.UNDEFINED
@pytest.mark.unit_test
def test_phase_context_double_nested():
graph_manager = GraphManager(name='', schedule_params=ScheduleParameters(), vis_params=VisualizationParameters())
assert graph_manager.phase == RunPhase.UNDEFINED
with graph_manager.phase_context(RunPhase.TRAIN):
assert graph_manager.phase == RunPhase.TRAIN
with graph_manager.phase_context(RunPhase.TEST):
assert graph_manager.phase == RunPhase.TEST
with graph_manager.phase_context(RunPhase.UNDEFINED):
assert graph_manager.phase == RunPhase.UNDEFINED
assert graph_manager.phase == RunPhase.TEST
assert graph_manager.phase == RunPhase.TRAIN
assert graph_manager.phase == RunPhase.UNDEFINED