"""
Tests for LearningEngine.
"""

import pytest
from dcos.learning.engine import LearningEngine, Experience


class TestLearningEngine:
    def test_record_experience(self):
        engine = LearningEngine()
        engine.record_experience(Experience(state={}, action="move", reward=1.0))
        assert engine.experience_count() == 1

    def test_best_action(self):
        engine = LearningEngine()
        engine.record_experience(Experience(state={}, action="a", reward=10))
        engine.record_experience(Experience(state={}, action="b", reward=5))
        best = engine.best_action(["a", "b"])
        assert best == "a"

    def test_get_action_score(self):
        engine = LearningEngine()
        engine.record_experience(Experience(state={}, action="explore", reward=0.5))
        assert engine.get_action_score("explore") == 0.05  # 0.5 * 0.1