import unittest

import numpy as np

from tbase.common.cmd_util import set_global_seeds
from tbase.network.polices import RandomPolicy


class TestPolices(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        set_global_seeds(0)

    def test_random_policy(self):
        policy = RandomPolicy(2)
        # action 1
        actual = policy.select_action([])
        expected = [1.0, -0.2534131770209437]
        self.assertEqual(expected, list(actual.astype(np.float)))
        # action 2
        actual = policy.select_action([])
        expected = [-1.0, 0.8324962832376306]
        self.assertEqual(expected, list(actual.astype(np.float)))


if __name__ == '__main__':
    unittest.main()

[((27, 4, 27, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((12, 8, 12, 27), 'tbase.common.cmd_util.set_global_seeds', 'set_global_seeds', ({(12, 25, 12, 26): '(0)'}, {}), '(0)', False, 'from tbase.common.cmd_util import set_global_seeds\n'), ((15, 17, 15, 32), 'tbase.network.polices.RandomPolicy', 'RandomPolicy', ({(15, 30, 15, 31): '2'}, {}), '(2)', False, 'from tbase.network.polices import RandomPolicy\n')]