# Copyright 2022 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Iterable

import pytest
from _pytest.logging import LogCaptureFixture

from gpflow.experimental.check_shapes.config import (
    DocstringFormat,
    ShapeCheckingState,
    get_enable_check_shapes,
    get_enable_function_call_precompute,
    get_rewrite_docstrings,
    set_enable_check_shapes,
    set_enable_function_call_precompute,
    set_rewrite_docstrings,
)


@pytest.fixture(autouse=True)
def test_auto_graph_compile(caplog: LogCaptureFixture) -> Iterable[None]:
    yield

    for when in ["setup", "call", "teardown"]:
        for record in caplog.get_records(when):
            assert not record.msg.startswith("AutoGraph could not transform"), record.getMessage()


@pytest.fixture(autouse=True)
def enable_shape_checks() -> Iterable[None]:
    # Ensure that:
    # 1: `check_shapes` is enabled when running these tests.
    # 2: If a test manipulates `check_shapes` settings, they are reset after the test.
    # See also: tests/conftest.py
    old_enable = get_enable_check_shapes()
    old_rewrite_docstrings = get_rewrite_docstrings()
    old_function_call_precompute = get_enable_function_call_precompute()
    set_enable_check_shapes(ShapeCheckingState.ENABLED)
    set_rewrite_docstrings(DocstringFormat.SPHINX)
    set_enable_function_call_precompute(True)
    yield
    set_enable_function_call_precompute(old_function_call_precompute)
    set_rewrite_docstrings(old_rewrite_docstrings)
    set_enable_check_shapes(old_enable)