# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# 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.

"""Tests that are common to both st.cache_data and st.cache_resource"""

from __future__ import annotations

import threading
import time
import unittest
from datetime import timedelta
from typing import Any
from unittest.mock import MagicMock, Mock, patch

import pytest
from parameterized import parameterized

import streamlit as st
from streamlit.errors import StreamlitAPIException
from streamlit.runtime import Runtime
from streamlit.runtime.caching import (
    cache_data,
    cache_data_api,
    cache_resource,
    cache_resource_api,
    clear_session_data_cache,
    clear_session_resource_cache,
)
from streamlit.runtime.caching.cache_errors import CacheReplayClosureError
from streamlit.runtime.caching.cache_utils import CachedResult
from streamlit.runtime.caching.storage.dummy_cache_storage import (
    MemoryCacheStorageManager,
)
from streamlit.runtime.forward_msg_queue import ForwardMsgQueue
from streamlit.runtime.fragment import MemoryFragmentStorage
from streamlit.runtime.memory_uploaded_file_manager import MemoryUploadedFileManager
from streamlit.runtime.pages_manager import PagesManager
from streamlit.runtime.scriptrunner import (
    ScriptRunContext,
    add_script_run_ctx,
    get_script_run_ctx,
)
from streamlit.runtime.scriptrunner_utils import script_run_context
from streamlit.runtime.state import SafeSessionState, SessionState
from streamlit.testing.v1.app_test import AppTest
from tests.delta_generator_test_case import DeltaGeneratorTestCase
from tests.exception_capturing_thread import call_on_threads
from tests.streamlit.elements.image_test import create_image
from tests.testutil import create_mock_script_run_ctx


def get_text_or_block(delta):
    if delta.WhichOneof("type") == "new_element":
        element = delta.new_element
        if element.WhichOneof("type") == "text":
            return element.text.body
        return None
    if delta.WhichOneof("type") == "add_block":
        return "new_block"
    return None


def as_cached_result(value: Any) -> CachedResult:
    """Creates cached results for a function that returned `value`
    and did not execute any elements.
    """
    return CachedResult(value, [], st._main.id, st.sidebar.id)


class CommonCacheTest(DeltaGeneratorTestCase):
    def tearDown(self):
        # Clear caches
        st.cache_data.clear()
        st.cache_resource.clear()

        # And some tests create widgets, and can result in DuplicateWidgetID
        # errors on subsequent runs.
        ctx = script_run_context.get_script_run_ctx()
        if ctx is not None:
            ctx.widget_ids_this_run.clear()
            ctx.widget_user_keys_this_run.clear()

        super().tearDown()

    def get_text_delta_contents(self) -> list[str]:
        deltas = self.get_all_deltas_from_queue()
        return [
            element.text.body
            for element in (delta.new_element for delta in deltas)
            if element.WhichOneof("type") == "text"
        ]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_simple(self, _, cache_decorator):
        @cache_decorator
        def foo():
            return 42

        assert foo() == 42
        assert foo() == 42

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_multiple_int_like_floats(self, _, cache_decorator):
        @cache_decorator
        def foo(x):
            return x

        assert foo(1.0) == 1.0
        assert foo(3.0) == 3.0

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_return_cached_object(self, _, cache_decorator):
        """If data has been cached, the cache function shouldn't be called."""
        with patch.object(st, "exception") as mock_exception:
            called = [False]

            @cache_decorator
            def f(x):
                called[0] = True
                return x

            assert not called[0]
            f(0)

            assert called[0]

            called = [False]  # Reset called

            f(0)
            assert not called[0]

            f(1)
            assert called[0]

            mock_exception.assert_not_called()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_mutate_args(self, _, cache_decorator):
        """Mutating an argument inside a cached function doesn't throw
        an error (but it's probably not a great idea)."""
        with patch.object(st, "exception") as mock_exception:

            @cache_decorator
            def foo(d):
                d["answer"] += 1
                return d["answer"]

            d = {"answer": 0}

            assert foo(d) == 1
            assert foo(d) == 2

            mock_exception.assert_not_called()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_ignored_args(self, _, cache_decorator):
        """Args prefixed with _ are not used as part of the cache key."""
        call_count = [0]

        @cache_decorator
        def foo(arg1, _arg2, *args, kwarg1, _kwarg2=None, **kwargs):
            call_count[0] += 1

        foo(1, 2, 3, kwarg1=4, _kwarg2=5, kwarg3=6, _kwarg4=7)
        assert call_count == [1]

        # Call foo again, but change the values for _arg2, _kwarg2, and _kwarg4.
        # The call count shouldn't change, because these args will not be part
        # of the hash.
        foo(1, None, 3, kwarg1=4, _kwarg2=None, kwarg3=6, _kwarg4=None)
        assert call_count == [1]

        # Changing the value of any other argument will increase the call
        # count. We test each argument type:

        # > arg1 (POSITIONAL_OR_KEYWORD)
        foo(None, 2, 3, kwarg1=4, _kwarg2=5, kwarg3=6, _kwarg4=7)
        assert call_count == [2]

        # > *arg (VAR_POSITIONAL)
        foo(1, 2, None, kwarg1=4, _kwarg2=5, kwarg3=6, _kwarg4=7)
        assert call_count == [3]

        # > kwarg1 (KEYWORD_ONLY)
        foo(1, 2, 3, kwarg1=None, _kwarg2=5, kwarg3=6, _kwarg4=7)
        assert call_count == [4]

        # > **kwarg (VAR_KEYWORD)
        foo(1, 2, 3, kwarg1=4, _kwarg2=5, kwarg3=None, _kwarg4=7)
        assert call_count == [5]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_member_function(self, _, cache_decorator):
        """Our cache decorators can be applied to class member functions."""

        class TestClass:
            @cache_decorator
            def member_func(_self):
                # We underscore-prefix `_self`, because our class is not hashable.
                return "member func!"

            @classmethod
            @cache_decorator
            def class_method(cls):
                return "class method!"

            @staticmethod
            @cache_decorator
            def static_method():
                return "static method!"

        obj = TestClass()
        assert obj.member_func() == "member func!"
        assert obj.class_method() == "class method!"
        assert obj.static_method() == "static method!"

    @parameterized.expand(
        [
            ("cache_data", cache_data),
            ("cache_resource", cache_resource),
        ]
    )
    def test_cached_st_function_warning(self, _, cache_decorator):
        """Ensure we properly warn when interactive st.foo functions are called
        inside a cached function.
        """
        forward_msg_queue = ForwardMsgQueue()
        orig_report_ctx = get_script_run_ctx()
        add_script_run_ctx(
            threading.current_thread(),
            ScriptRunContext(
                session_id="test session id",
                _enqueue=forward_msg_queue.enqueue,
                query_string="",
                session_state=SafeSessionState(SessionState(), lambda: None),
                uploaded_file_mgr=MemoryUploadedFileManager("/mock/upload"),
                main_script_path="",
                user_info={"email": "test@example.com"},
                fragment_storage=MemoryFragmentStorage(),
                pages_manager=PagesManager(""),
            ),
        )

        with patch.object(st, "exception") as warning:
            st.text("foo")
            warning.assert_not_called()

            @cache_decorator
            def cached_func():
                st.text("Inside cached func")

            cached_func()
            warning.assert_not_called()

            warning.reset_mock()

            # Make sure everything got reset properly
            st.text("foo")
            warning.assert_not_called()

            # Test nested st.cache functions
            @cache_decorator
            def outer():
                @cache_decorator
                def inner():
                    st.text("Inside nested cached func")

                return inner()

            outer()
            warning.assert_not_called()

            warning.reset_mock()

            # Test st.cache functions that raise errors
            with pytest.raises(RuntimeError):

                @cache_decorator
                def cached_raise_error():
                    st.text("About to throw")
                    raise RuntimeError("avast!")

                cached_raise_error()

            warning.assert_not_called()
            warning.reset_mock()

            # Make sure everything got reset properly
            st.text("foo")
            warning.assert_not_called()

            # Test st.cache functions with widgets
            @cache_decorator
            def cached_widget():
                st.button("Press me!")

            cached_widget()

            warning.assert_called()
            warning.reset_mock()

            # Make sure everything got reset properly
            st.text("foo")
            warning.assert_not_called()

            add_script_run_ctx(threading.current_thread(), orig_report_ctx)

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay(self, _, cache_decorator):
        @cache_decorator
        def foo_replay(i):
            st.text(i)
            return i

        foo_replay(1)
        st.text("---")
        foo_replay(1)

        text = self.get_text_delta_contents()

        assert text == ["1", "---", "1"]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_nested(self, _, cache_decorator):
        @cache_decorator
        def inner(i):
            st.text(i)

        @cache_decorator
        def outer(i):
            inner(i)
            st.text(i + 10)

        outer(1)
        outer(1)
        st.text("---")
        inner(2)
        outer(2)
        st.text("---")
        outer(3)
        inner(3)

        text = self.get_text_delta_contents()
        assert text == [
            "1",
            "11",
            "1",
            "11",
            "---",
            "2",
            "2",
            "12",
            "---",
            "3",
            "13",
            "3",
        ]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_outer_blocks(self, _, cache_decorator):
        @cache_decorator
        def foo(i):
            st.text(i)
            return i

        with st.container():
            foo(1)
            st.text("---")
            foo(1)

        text = self.get_text_delta_contents()
        assert text == ["1", "---", "1"]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_sidebar(self, _, cache_decorator):
        @cache_decorator(show_spinner=False)
        def foo(i):
            st.sidebar.text(i)
            return i

        foo(1)  # [1,0]
        st.text("---")  # [0,0]
        foo(1)  # [1,1]

        text = [
            get_text_or_block(delta)
            for delta in self.get_all_deltas_from_queue()
            if get_text_or_block(delta) is not None
        ]
        assert text == ["1", "---", "1"]

        paths = [
            msg.metadata.delta_path
            for msg in self.forward_msg_queue._queue
            if msg.HasField("delta")
        ]
        assert paths == [[1, 0], [0, 0], [1, 1]]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_inner_blocks(self, _, cache_decorator):
        @cache_decorator(show_spinner=False)
        def foo(i):
            with st.container():
                st.text(i)
                return i

        with st.container():  # [0,0]
            st.text(0)  # [0,0,0]
        st.text("---")  # [0,1]
        with st.container():  # [0,2]
            st.text(0)  # [0,2,0]

        foo(1)  # [0,3] and [0,3,0]
        st.text("---")  # [0,4]
        foo(1)  # [0,5] and [0,5,0]

        paths = [
            msg.metadata.delta_path
            for msg in self.forward_msg_queue._queue
            if msg.HasField("delta")
        ]
        assert paths == [
            [0, 0],
            [0, 0, 0],
            [0, 1],
            [0, 2],
            [0, 2, 0],
            [0, 3],
            [0, 3, 0],
            [0, 4],
            [0, 5],
            [0, 5, 0],
        ]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_inner_direct(self, _, cache_decorator):
        @cache_decorator(show_spinner=False)
        def foo(i):
            cont = st.container()
            cont.text(i)
            return i

        foo(1)  # [0,0] and [0,0,0]
        st.text("---")  # [0,1]
        foo(1)  # [0,2] and [0,2,0]

        text = self.get_text_delta_contents()
        assert text == ["1", "---", "1"]

        paths = [
            msg.metadata.delta_path
            for msg in self.forward_msg_queue._queue
            if msg.HasField("delta")
        ]
        assert paths == [[0, 0], [0, 0, 0], [0, 1], [0, 2], [0, 2, 0]]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_function_replay_outer_direct(self, _, cache_decorator):
        cont = st.container()

        @cache_decorator
        def foo(i):
            cont.text(i)
            return i

        with pytest.raises(CacheReplayClosureError):
            foo(1)
            st.text("---")
            foo(1)

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_cached_st_image_replay(self, _, cache_decorator):
        """Basic sanity check that nothing blows up. This test assumes that
        actual caching/replay functionality are covered by e2e tests that
        can more easily test them.
        """

        @cache_decorator
        def img_fn():
            st.image(create_image(10))

        img_fn()
        img_fn()

        @cache_decorator
        def img_fn_multi():
            st.image([create_image(5), create_image(15), create_image(100)])

        img_fn_multi()
        img_fn_multi()

    @parameterized.expand(
        [
            ("cache_data", cache_data, cache_data.clear),
            ("cache_resource", cache_resource, cache_resource.clear),
        ]
    )
    def test_clear_all_caches(self, _, cache_decorator, clear_cache_func):
        """Calling a cache's global `clear_all` function should remove all
        items from all caches of the appropriate type.
        """
        foo_vals = []

        @cache_decorator
        def foo(x):
            foo_vals.append(x)
            return x

        bar_vals = []

        @cache_decorator
        def bar(x):
            bar_vals.append(x)
            return x

        foo(0), foo(1), foo(2)
        bar(0), bar(1), bar(2)
        assert foo_vals == [0, 1, 2]
        assert bar_vals == [0, 1, 2]

        # Clear the cache and access our original values again. They
        # should be recomputed.
        clear_cache_func()

        foo(0), foo(1), foo(2)
        bar(0), bar(1), bar(2)
        assert foo_vals == [0, 1, 2, 0, 1, 2]
        assert bar_vals == [0, 1, 2, 0, 1, 2]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_clear_single_cache(self, _, cache_decorator):
        foo_call_count = [0]

        @cache_decorator
        def foo():
            foo_call_count[0] += 1

        bar_call_count = [0]

        @cache_decorator
        def bar():
            bar_call_count[0] += 1

        foo(), foo(), foo()
        bar(), bar(), bar()
        assert foo_call_count[0] == 1
        assert bar_call_count[0] == 1

        # Clear just foo's cache, and call the functions again.
        foo.clear()

        foo(), foo(), foo()
        bar(), bar(), bar()

        # Foo will have been called a second time, and bar will still
        # have been called just once.
        assert foo_call_count[0] == 2
        assert bar_call_count[0] == 1

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_without_spinner(self, _, cache_decorator):
        """If the show_spinner flag is not set, the report queue should be
        empty.
        """

        @cache_decorator(show_spinner=False)
        def function_without_spinner(x: int) -> int:
            return x

        function_without_spinner(3)
        assert self.forward_msg_queue.is_empty()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_with_spinner(self, _, cache_decorator):
        """If the show_spinner flag is set, there should be one element in the
        report queue.
        """

        @cache_decorator(show_spinner=True)
        def function_with_spinner(x: int) -> int:
            return x

        function_with_spinner(3)
        assert not self.forward_msg_queue.is_empty()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_with_custom_text_spinner(self, _, cache_decorator):
        """If the show_spinner flag is set, there should be one element in the
        report queue.
        """

        @cache_decorator(show_spinner="CUSTOM_TEXT")
        def function_with_spinner_custom_text(x: int) -> int:
            return x

        function_with_spinner_custom_text(3)
        assert not self.forward_msg_queue.is_empty()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_with_empty_text_spinner(self, _, cache_decorator):
        """If the show_spinner flag is set, even if it is empty text,
        there should be one element in the report queue.
        """

        @cache_decorator(show_spinner="")
        def function_with_spinner_empty_text(x: int) -> int:
            return x

        function_with_spinner_empty_text(3)
        assert not self.forward_msg_queue.is_empty()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_spinner_with_nested_cached_functions(self, _, cache_decorator):
        """If a cached function calls another cached function, only one spinner
        should be created.
        """

        @cache_decorator(show_spinner="")
        def inner(x: int) -> int:
            return x

        @cache_decorator(show_spinner="")
        def outer(x: int) -> int:
            return inner(x)

        outer(3)
        assert not self.forward_msg_queue.is_empty()

        # The spinner uses a transient element and shows the spinner only
        # after a timeout. Instead of mocking the time and waiting for the
        # timeout, we check for the transient element with an empty set
        # of elements. This represents the spinner element disappearing.
        transient_elements_count = 0
        for msg in self.forward_msg_queue._queue:
            if (
                msg.HasField("delta")
                and msg.delta.HasField("new_transient")
                and len(msg.delta.new_transient.elements) == 0
            ):
                transient_elements_count += 1
        # Since we automatically prevent spinners for nested cached functions,
        # there should only be a single transient spinner element.
        assert transient_elements_count == 1

    @parameterized.expand(
        [
            ("cache_data", cache_data, cache_data_api),
            (
                "cache_resource",
                cache_resource,
                cache_resource_api,
            ),
        ]
    )
    def test_session_scope_handles_lookup(self, _, cache_decorator, cache_module):
        """A cache should work with session scoping."""

        counter = 0
        curr_session_id: str

        @cache_decorator(scope="session")
        def get_value(val: str) -> str:
            nonlocal counter, curr_session_id
            counter += 1
            return f"{val}-{curr_session_id}-{counter}"

        # Simulate one session.
        curr_session_id = "aaa"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            # Validate function is called twice for new args.
            result_a_first = get_value("first")
            assert result_a_first == "first-aaa-1"
            result_a_second = get_value("second")
            assert result_a_second == "second-aaa-2"

            # Validate function is not called again for same args.
            assert result_a_first == get_value("first")
            assert result_a_second == get_value("second")

        # Simulate a second session.
        curr_session_id = "bbbb"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            # Validate function is called twice for new args.
            result_b_first = get_value("first")
            assert result_b_first == "first-bbbb-3"
            result_b_second = get_value("second")
            assert result_b_second == "second-bbbb-4"

            # Validate function is not called again for same args.
            assert result_b_first == get_value("first")
            assert result_b_second == get_value("second")

        # Validate the first session still has a cache.
        # Simulate one session.
        curr_session_id = "aaa"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            assert result_a_first == get_value("first")
            assert result_a_second == get_value("second")

    @parameterized.expand(
        [
            ("cache_data", cache_data, cache_data_api, clear_session_data_cache),
            (
                "cache_resource",
                cache_resource,
                cache_resource_api,
                clear_session_resource_cache,
            ),
        ]
    )
    def test_session_scope_handles_clear(
        self, _, cache_decorator, cache_module, clear_session_cache
    ):
        """Clearing a single session's scope should work."""
        counter = 0
        curr_session_id: str

        @cache_decorator(scope="session")
        def get_value(val: str) -> str:
            nonlocal counter, curr_session_id
            counter += 1
            return f"{val}-{curr_session_id}-{counter}"

        # Simulate one session.
        curr_session_id = "aaa"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            assert get_value("first") == "first-aaa-1"
            assert get_value("second") == "second-aaa-2"

        # Simulate a second session.
        curr_session_id = "bbbb"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            result_b_first = get_value("first")
            assert result_b_first == "first-bbbb-3"
            result_b_second = get_value("second")
            assert result_b_second == "second-bbbb-4"

        # Clear the first session, and assert we get new values.
        curr_session_id = "aaa"
        clear_session_cache(curr_session_id)
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            assert get_value("first") == "first-aaa-5"
            assert get_value("second") == "second-aaa-6"

        # Validate that the second session remained unchanged.
        curr_session_id = "bbbb"
        with patch.object(
            cache_module, "get_session_id_or_throw"
        ) as mock_get_session_id:
            mock_get_session_id.return_value = curr_session_id

            assert result_b_first == get_value("first")
            assert result_b_second == get_value("second")

    @parameterized.expand(
        [
            ("cache_data", cache_data),
            ("cache_resource", cache_resource),
        ]
    )
    def test_bad_scope_raises_exception(self, _, cache_decorator):
        """A bad scope argument should raise an exception."""

        with pytest.raises(StreamlitAPIException) as e:

            @cache_decorator(scope="request")
            def get_foo() -> None:
                pass

        e.match("Unsupported scope option.*request")


class CommonCacheTTLTest(unittest.TestCase):
    def setUp(self) -> None:
        # Caching functions rely on an active script run ctx
        add_script_run_ctx(threading.current_thread(), create_mock_script_run_ctx())
        mock_runtime = MagicMock(spec=Runtime)
        mock_runtime.cache_storage_manager = MemoryCacheStorageManager()
        Runtime._instance = mock_runtime

    def tearDown(self):
        cache_data.clear()
        cache_resource.clear()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    @patch("streamlit.runtime.caching.cache_utils.TTLCACHE_TIMER")
    def test_ttl(self, _, cache_decorator, timer_patch: Mock):
        """Entries should expire after the given ttl."""
        one_day = 60 * 60 * 24

        # Create 2 cached functions to test that they don't interfere
        # with each other.
        foo_vals = []

        @cache_decorator(ttl=one_day)
        def foo(x):
            foo_vals.append(x)
            return x

        bar_vals = []

        @cache_decorator(ttl=one_day * 2)
        def bar(x):
            bar_vals.append(x)
            return x

        # Store a value at time 0
        timer_patch.return_value = 0
        foo(0)
        bar(0)
        assert foo_vals == [0]
        assert bar_vals == [0]

        # Advance our timer, but not enough to expire our value.
        timer_patch.return_value = one_day * 0.5
        foo(0)
        bar(0)
        assert foo_vals == [0]
        assert bar_vals == [0]

        # Advance our timer enough to expire foo, but not bar.
        timer_patch.return_value = one_day * 1.5
        foo(0)
        bar(0)
        assert foo_vals == [0, 0]
        assert bar_vals == [0]

        # Expire bar. Foo's second value was inserted at time=1.5 days,
        # so it won't expire until time=2.5 days
        timer_patch.return_value = (one_day * 2) + 1
        foo(0)
        bar(0)
        assert foo_vals == [0, 0]
        assert bar_vals == [0, 0]

        # Expire foo for a second time.
        timer_patch.return_value = (one_day * 2.5) + 1
        foo(0)
        bar(0)
        assert foo_vals == [0, 0, 0]
        assert bar_vals == [0, 0]

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    @patch("streamlit.runtime.caching.cache_utils.TTLCACHE_TIMER")
    def test_ttl_timedelta(self, _, cache_decorator, timer_patch: Mock):
        """Entries should expire after the given ttl."""
        one_day_seconds = 60 * 60 * 24
        one_day_timedelta = timedelta(days=1)
        two_days_timedelta = timedelta(days=2)

        # Create 2 cached functions to test that they don't interfere
        # with each other.
        foo_vals = []

        @cache_decorator(ttl=one_day_timedelta)
        def foo(x):
            foo_vals.append(x)
            return x

        bar_vals = []

        @cache_decorator(ttl=two_days_timedelta)
        def bar(x):
            bar_vals.append(x)
            return x

        # Store a value at time 0
        timer_patch.return_value = 0
        foo(0)
        bar(0)
        assert foo_vals == [0]
        assert bar_vals == [0]

        # Advance our timer, but not enough to expire our value.
        timer_patch.return_value = one_day_seconds * 0.5
        foo(0)
        bar(0)
        assert foo_vals == [0]
        assert bar_vals == [0]

        # Advance our timer enough to expire foo, but not bar.
        timer_patch.return_value = one_day_seconds * 1.5
        foo(0)
        bar(0)
        assert foo_vals == [0, 0]
        assert bar_vals == [0]

        # Expire bar. Foo's second value was inserted at time=1.5 days,
        # so it won't expire until time=2.5 days
        timer_patch.return_value = (one_day_seconds * 2) + 1
        foo(0)
        bar(0)
        assert foo_vals == [0, 0]
        assert bar_vals == [0, 0]

        # Expire foo for a second time.
        timer_patch.return_value = (one_day_seconds * 2.5) + 1
        foo(0)
        bar(0)
        assert foo_vals == [0, 0, 0]
        assert bar_vals == [0, 0]


class CommonCacheThreadingTest(unittest.TestCase):
    # The number of threads to run our tests on
    NUM_THREADS = 50

    def setUp(self):
        mock_runtime = MagicMock(spec=Runtime)
        mock_runtime.cache_storage_manager = MemoryCacheStorageManager()
        Runtime._instance = mock_runtime

    def tearDown(self):
        # Some of these tests reach directly into CALL_STACK data and twiddle it.
        # Reset default values on teardown.

        # Clear caches
        st.cache_data.clear()
        st.cache_resource.clear()

        # And some tests create widgets, and can result in DuplicateWidgetID
        # errors on subsequent runs.
        ctx = script_run_context.get_script_run_ctx()
        if ctx is not None:
            ctx.widget_ids_this_run.clear()
            ctx.widget_user_keys_this_run.clear()

        super().tearDown()

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_get_cache(self, _, cache_decorator):
        """Accessing a cached value is safe from multiple threads."""

        cached_func_call_count = [0]

        @cache_decorator
        def foo():
            cached_func_call_count[0] += 1
            return 42

        def call_foo(_: int) -> None:
            assert foo() == 42

        # Call foo from multiple threads and assert no errors.
        call_on_threads(call_foo, self.NUM_THREADS)

        # The cached function should only be called once (see `test_compute_value_only_once`).
        assert cached_func_call_count[0] == 1

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_compute_value_only_once(self, _, cache_decorator):
        """Cached values should be computed only once, even if multiple sessions read from an
        unwarmed cache simultaneously.
        """
        cached_func_call_count = [0]

        @cache_decorator
        def foo():
            assert cached_func_call_count[0] == 0, (
                "A cached value was computed multiple times!"
            )
            cached_func_call_count[0] += 1

            # Sleep to "guarantee" that our other threads try to access the
            # cached data while it's being computed. (The other threads should
            # block on cache computation, so this function should only
            # be called a single time.)
            time.sleep(0.25)
            return 42

        def call_foo(_: int) -> None:
            assert foo() == 42

        call_on_threads(call_foo, num_threads=self.NUM_THREADS, timeout=0.5)

    @parameterized.expand(
        [
            ("cache_data", cache_data, cache_data.clear),
            ("cache_resource", cache_resource, cache_resource.clear),
        ]
    )
    def test_clear_all_caches(self, _, cache_decorator, clear_cache_func):
        """Clearing all caches is safe to call from multiple threads."""

        @cache_decorator
        def foo():
            return 42

        # Populate the cache
        foo()

        def clear_caches(_: int) -> None:
            clear_cache_func()

        # Clear the cache from a bunch of threads and assert no errors.
        call_on_threads(clear_caches, self.NUM_THREADS)

        # Sanity check: ensure we can still call our cached function.
        assert foo() == 42

    @parameterized.expand(
        [("cache_data", cache_data), ("cache_resource", cache_resource)]
    )
    def test_clear_single_cache(self, _, cache_decorator):
        """It's safe to clear a single function cache from multiple threads."""

        @cache_decorator
        def foo():
            return 42

        # Populate the cache
        foo()

        def clear_foo(_: int) -> None:
            foo.clear()

        # Clear it from a bunch of threads and assert no errors.
        call_on_threads(clear_foo, self.NUM_THREADS)

        # Sanity check: ensure we can still call our cached function.
        assert foo() == 42


def test_arrow_replay():
    """Regression test for https://github.com/streamlit/streamlit/issues/6103"""
    at = AppTest.from_file("test_data/arrow_replay.py").run()

    assert not at.exception
