from unittest.mock import Mock, patch

import numpy as np

from game.models import ValuePolicyModel


def test_predict():
    mask = np.zeros((9, 7, 8), dtype=bool)
    mask[1, 2, 3] = 1
    mask[6, 6, 6] = 1

    tensor_mock = Mock()

    policy_tensor = np.zeros((9, 7, 8), dtype=float)
    policy_tensor[0, 0, 0] = 10
    policy_tensor[1, 2, 3] = 100
    policy_tensor[6, 6, 6] = 100
    policy_tensor = policy_tensor.reshape(-1)

    value = np.array([[0.7]], dtype=float)

    get_prediction_mock = Mock(return_value=(value, policy_tensor))
    network_mock = Mock(spec=ValuePolicyModel, output_shape=(9, 7, 8),
                        input_shape=(9, 7, 178), _get_prediction=get_prediction_mock)

    result_value, result_policy = ValuePolicyModel.predict(
        self=network_mock, tensor=tensor_mock, mask=mask)

    get_prediction_mock.assert_called_once_with(tensor_mock)

    expected_value = 0.7

    expected_policy = np.zeros((9, 7, 8))
    expected_policy[1, 2, 3] = 0.5
    expected_policy[6, 6, 6] = 0.5

    assert isinstance(result_value, float)
    assert result_value == expected_value

    assert np.array_equal(result_policy, expected_policy)

[(9, 'numpy.zeros', 'np.zeros', (['(9, 7, 8)'], {'dtype': 'bool'}), True, 'import numpy as np\n'), (13, 'unittest.mock.Mock', 'Mock', ([], {}), False, 'from unittest.mock import Mock, patch\n'), (15, 'numpy.zeros', 'np.zeros', (['(9, 7, 8)'], {'dtype': 'float'}), True, 'import numpy as np\n'), (21, 'numpy.array', 'np.array', (['[[0.7]]'], {'dtype': 'float'}), True, 'import numpy as np\n'), (23, 'unittest.mock.Mock', 'Mock', ([], {'return_value': '(value, policy_tensor)'}), False, 'from unittest.mock import Mock, patch\n'), (24, 'unittest.mock.Mock', 'Mock', ([], {'spec': 'ValuePolicyModel', 'output_shape': '(9, 7, 8)', 'input_shape': '(9, 7, 178)', '_get_prediction': 'get_prediction_mock'}), False, 'from unittest.mock import Mock, patch\n'), (27, 'game.models.ValuePolicyModel.predict', 'ValuePolicyModel.predict', ([], {'self': 'network_mock', 'tensor': 'tensor_mock', 'mask': 'mask'}), False, 'from game.models import ValuePolicyModel\n'), (34, 'numpy.zeros', 'np.zeros', (['(9, 7, 8)'], {}), True, 'import numpy as np\n'), (41, 'numpy.array_equal', 'np.array_equal', (['result_policy', 'expected_policy'], {}), True, 'import numpy as np\n')]