Streamlit Unit Testing - infinite loop with st.rerun

I have recently come across Streamlits (1.28.2) Unit Testing and attempted to give it a go - as this is a feature I have been hoping will come for a while!

However this is my first experience with unit tests and do lack much of the basic knowledge - because of that i am hoping you can help me with a problem i am having

I seem to be getting stuck in an infinite loop when attempting to test a button click in my app - the button click is designed to take in a text input and write that as commentary to a connected DB (which will be displayed in a series of other streamlit apps)

Here is the code I am using to test the app along with the snippet of code which runs once the button is pressed (inside the app script)

def test_com_input():
    at = AppTest.from_file('1_metric_inputs.py').run(timeout = 600)

    at.text_area[0].input('This is a test from the unit tests pt2').run(timeout = 60)

    at.button[0].click().run(timeout = 15)

    assert not at.exception
clicked = st.form_submit_button(label="Submit")

if clicked: 
    passed = inputs.submit_coms(st.session_state['coms'], sdi_id, month_str, com_lim, user)

    if passed:
        st.success('Commentary Updated!')

        time.sleep(2)

        st.rerun()

I believe the problem is the st.rerun - which is causing the following error as well as the commentary to be written to the DB over and over again

FAILED src/test_simple.py::test_com_input - KeyError: 'client_state'

Any advice on how i could remedy this/best practice for unit testing would be much appreciated

Cheers
Toby

That error message is unrelated to the code you posted.

Do you know what the error message is related to/ how to solve it?

I posted that code - because if i comment out the st.rerun() from the app - then the test runs exactly as expected and passes - but we need the st.rerun() for the app to work as intended

It is related to a test called test_simple.py.

That is the name of the file which then has the test_com_input() function inside - which is the code snippet i sent (top input)

You are right, I missed that. But then the message says that 'client_state'
is being used as a key, and I can’t find anything like that in your code.

Yes - i believe the error is being raised from the Streamlit testing framework, not any code I wrote. but I think the problem is to do with the st.rerun line in my app as I seem to get stuck in an infinite loop - I can share the code where the error is being raised from - but as I say this is not code I wrote

self = AppTest(_script_path='/home/gaskeltx3/projects/cdi-streamlit/src/1_metric_inputs.py', default_timeout=3, session_state...', label='Submit', form_id='value'), 3: Success()})}), 1: SpecialBlock(type='sidebar'), 2: SpecialBlock(type='event')})
widget_state = widgets {
  id: "$$WIDGET_ID-5f7fd2dc05675ae6e976eef687a15d2b-None"
  int_value: 0
}
widgets {
  id: "$$WIDGET_ID-46bf...3"
}
widgets {
  id: "$$WIDGET_ID-2345ae15013aad269664939a78d89534-FormSubmitter:value-Submit"
  trigger_value: true
}

timeout = 15

    def _run(
        self,
        widget_state: WidgetStates | None = None,
        timeout: float | None = None,
    ) -> AppTest:
        """Run the script, and parse the output messages for querying
        and interaction.

        Timeout is in seconds, or None to use the default timeout of the runner.
        """
        # Have to import the streamlit module itself so replacing st.secrets
        # is visible to other modules.
        import streamlit as st

        if timeout is None:
            timeout = self.default_timeout

        # setup
        mock_runtime = MagicMock(spec=Runtime)
        mock_runtime.media_file_mgr = MediaFileManager(
            MemoryMediaFileStorage("/mock/media")
        )
        mock_runtime.cache_storage_manager = MemoryCacheStorageManager()
        Runtime._instance = mock_runtime
        with source_util._pages_cache_lock:
            saved_cached_pages = source_util._cached_pages
            source_util._cached_pages = None

        saved_secrets: Secrets = st.secrets
        # Only modify global secrets stuff if we have been given secrets
        if self.secrets:
            new_secrets = Secrets([])
            new_secrets._secrets = self.secrets
            st.secrets = new_secrets

        script_runner = LocalScriptRunner(self._script_path, self.session_state)
        self._tree = script_runner.run(widget_state, self.query_params, timeout)
        self._tree._runner = self
        # Last event is SHUTDOWN, so the corresponding data includes query string
>       query_string = script_runner.event_data[-1]["client_state"].query_string
E       KeyError: 'client_state'


It looks like a case where the AppTest instance behaves differently than actually running the app. This will print CLICKED! forever and at some point it will also print the KeyError: 'client_state' exception (only once).

from streamlit.testing.v1 import AppTest

code = """
import time

import streamlit as st

if st.button(label="Submit"):
    print("CLICKED!")
    time.sleep(1)
    st.rerun()
"""

at = AppTest.from_string(code).run()
at.button[0].click().run()

However running the application normally it prints CLICKED! once per click, and no exceptions.

Note that calling time.sleep() is not necessary to trigger the bug but it makes the output more comfortable to look at.

There was s similar bug affecting normal execution of apps, it was fixed quickly but maybe the fix must be applied separately to streamlit.testing.

I created a new issue for this.

3 Likes