Mocking out a postgres connection in session state with pytest/monkeypatch

taylor1 · November 18, 2024, 3:23pm

Running locally, Python 3.12, Streamlit 1.40.1

Hi, please forgive the big setup here, trying to map my problem exactly. I am writing a bunch of tests for my streamlit app, and so far it has worked well using standard pytest and mocking out my Postgres queries with monkeypatch. However, I’m now trying to recover data from an INSERT ... RETURNING ...; query and having a lot of trouble getting the scoping on my mock correct.

Imagine I have an app setup as follows:

.
├── poetry.lock
├── pyproject.toml
└── streamlit
    ├── app.py
    ├── classes
    │   ├── __init__.py
    │   └── burg.py
    └── tests
        ├── __init__.py
        └── test_app.py

Where my main app is in app.py, and I have a class doing the lifting in classes/burg.py.

burg.py

import streamlit as st
from sqlalchemy.sql import text


class Burg:
    pass


class Burger(Burg):
    @staticmethod
    def flip_burger() -> list[str]:
        with st.session_state.conn.session as session:
            new_antigens = session.execute(
                text(
                    "INSERT INTO user_logs (update_time, antigens_id, users_id, "
                    "site_page, status_changed) VALUES ('2024-10-10 00:00:00', 1, 1, "
                    "'burgers', 'flipped') returning user_logs_id;"
                )
            )
            session.commit()
        new_ids = [str(new_id[0]) for new_id in new_antigens]
        return new_ids

app.py

import streamlit as st
from classes.burg import Burger

st.session_state.conn = st.connection("postgres")  # init psql connection


Burger.flip_burger()

test_app.py

import streamlit as st
import pandas as pd
from classes.burg import Burger


class Test_Burger:

    def test_flip_burger(self, monkeypatch):
        class session_state:
            class conn:
                class session:
                    def __enter__(self):
                        return self

                    def __exit__(self):
                        pass

                    @staticmethod
                    def execute(null):
                        return pd.DataFrame({"ids": [100]})

        def mock_conn(*args, **kwargs):
            return session_state()

        monkeypatch.setattr(st, "session_state", mock_conn)
        borger = Burger()
        assert borger.flip_burger() == [100]

The gist is I am inserting into my logs table and this function returns the ids of the rows that have been inserted.

I would like to use the monkeypatched method above to mock st.session_state.conn.session and .execute() so I can returned a couple fake ids and test my function.

I’ve tried a few different iterations of the session_state mock class (from How to monkeypatch/mock modules and environments - pytest documentation), but the closest I can get is:

FAILED tests/test_app.py::Test_Burger::test_flip_burger - AttributeError: 'function' object has no attribute 'conn'

I have a feeling I’m just not scoping or designing my session_state replacement correctly. What would be the right way to mock this out with monkeypatch?

Goyo · November 18, 2024, 6:07pm

You must be running a different test_app.py. because the one you posted here does nothing but defining a class Test_Burger that is never used.

taylor1 · November 18, 2024, 6:12pm

Yup, my mistake! Whole thing is:

import streamlit as st
import pandas as pd
from classes.burg import Burger


class Test_Burger:

    def test_flip_burger(self, monkeypatch):
        class session_state:
            class conn:
                class session:
                    def __enter__(self):
                        return self

                    def __exit__(self):
                        pass

                    @staticmethod
                    def execute(null):
                        return pd.DataFrame({"ids": [100]})

        def mock_conn(*args, **kwargs):
            return session_state()

        monkeypatch.setattr(st, "session_state", mock_conn)
        borger = Burger()
        assert borger.flip_burger() == [100]

I’ll fix above, careful since the code block has a scroll bar.

Goyo · November 18, 2024, 6:17pm

Same thing. That code does noting but defining Test_Burger .

taylor1 · November 18, 2024, 6:34pm

There’s a part at the bottom that calls the class method in question

        def mock_conn(*args, **kwargs):
            return session_state()

        monkeypatch.setattr(st, "session_state", mock_conn)
        borger = Burger()
        assert borger.flip_burger() == [100]

If you’re referring to the assertion being inside the class, PyTest is able to find and run the tests.

Goyo · November 18, 2024, 7:17pm

So you are doing:

monkeypatch.setattr(st, "session_state", mock_conn)

which amounts to temporarily having

st.session_state = mock_conn

Then when Burger is instantiated, this expression is evaluated

st.session_state.conn.session

which is now the same as

mock_conn.conn.session

But mock_conn is a function and it has no attribute conn, so an exception is raised.

I don’t know how to fix it yet, because I am still trying to figure out what the intent of your test is.

Anyway, this assertion will always fail:

assert borger.flip_burger() == [100]

because flip_burger returns a list of strings. Do you mean this instead?

assert borger.flip_burger() == ["100"]

taylor1 · November 18, 2024, 7:56pm

The intent of the test is to mock out with st.session_state.conn.session as session:, more specifically session.execute() which returns ids into new_antigens.

I can change it to monkeypatch.setattr(st, "session_state", mock_conn()) which will return an instance of session_state but that has a problem with the context manager:
FAILED streamlit/tests/test_app.py::Test_Burger::test_flip_burger - TypeError: 'type' object does not support the context manager protocol

Goyo · November 18, 2024, 8:35pm

The purpose of tests is usually make assertions about the behavior of an object. Mocking out something for the sake of mocking out tells you nothing about the behavior of your objects. There must be some behavior you want to test?

Right now I can see that you are asserting that borger.flip_burger() == [100] , which can never happen because borger.flip_burger() returns a list of strings.

Now it can be that you actually mean this instead:

borger.flip_burger() == [100]

In that case, yes, you can put your application in a state such that the assertion succeed. But it depends on the behavior of an sqlalchemy session, and there is no way for you to manipulate the session other than monkeypatching, and monkeypatching is tricky, as you already discovered.

So when it comes to monkeypatching, my advice would be: don’t do it. Make the dependency explicit instead. I would change your Burger class to this:

class Burger:
    @staticmethod
    def flip_burger(session) -> list[str]:
        new_antigens = session.execute(
            text(
                "INSERT INTO user_logs (update_time, antigens_id, users_id, "
                "site_page, status_changed) VALUES ('2024-10-10 00:00:00', 1, 1, "
                "'burgers', 'flipped') returning user_logs_id;"
            )
        )
        session.commit()
        new_ids = [str(new_id[0]) for new_id in new_antigens]
        return new_ids

Actually I would get rid of Burger because it is just a function in disguise, but whatever…

Then in the test:

class FakeSession:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return None

    def execute(self, statement):
        return [[100]]

    def commit(self):
        pass


class Test_Burger:
    def test_flip_burger(self):
        borger = Burger()
        with FakeSession() as fake_session:
            assert borger.flip_burger(session=fake_session) == ["100"]

I don’t se any value in a test like this, so maybe this is not what you are trying to do, but it is all I can get from your code.

taylor1 · November 18, 2024, 9:03pm

I am interested in testing flip_burger()'s ability to execute an INSERT statement and parse the output. This is also a small example, I perform other similar insertions throughout my app, so having this testing ability would prove overall useful.

I’m mostly interested in mocking here since I can use it to replace the session object in a way that lets me keep my application code as-is.

The real question is, and apologies if I’ve buried the lede here, “how can I mock the context manager to deliver results akin to actually running the SQL command?” It’s not necessarily a streamlit problem, but I thought others might have experience with using streamlit’s session state in this way (storing a database connection).

Goyo · November 18, 2024, 11:39pm

You mean monkey-patching. Well, you can do that too, it just requires more set-up.

class FakeSession:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return None

    def execute(self, statement):
        return [[100]]

    def commit(self):
        pass


class FakeConnection:
    session = FakeSession()


class FakeSessionState:
    conn = FakeConnection()


class Test_Burger:
    def test_flip_burger(self, monkeypatch):
        fake_session_state = FakeSessionState()
        monkeypatch.setattr(st, "session_state", fake_session_state)
        borger = Burger()
        assert borger.flip_burger() == ["100"]

But note that you are not checking the INSERT statement here. You could pass any argument to session.execute() in flip_burger(), and the test would still pass.

taylor1 · November 19, 2024, 2:51pm

Incredible work, thank you so much! Apologies if I wasn’t clear from the get go, it felt like I had to describe a lot of moving parts to get to the core of the problem and this fake burger class was my simplest example case.

I realize I’m not testing the insert itself (that’s sqlalchemy’s job), all along I wanted to test that my method could parse the results of the insertion, that’s all.

Is there any easier away to do this kind of test than to monkeypatch the session like this? Obviously it’s kind of niche but I figure this isn’t the most uncommon use case.

Goyo · November 19, 2024, 7:21pm

Actually you (and I) had, compared to the simplicity of the behavior under test. I had to monkeypatch st.session_state with a test double, then make it have a conn attribute which is itself another test double, then make it have a session attribute which is another test double, just to have the method under test call a couple of methods on the session.

I don’t think so. Not without changing the code. The code is hard to test so the tests will be ugly, hard to write and hard to read.

You could use Mock objects instead of writing the doubles yourself, and contextlib to make the fake session a context manager, but I don’t think it would make it easier or more readable, so I didn’t even try.

taylor1 · November 19, 2024, 7:35pm

To me it’s not too bad, since the code is a nested jumble, it only makes sense for the test to reflect that.

Thanks again!

system · November 21, 2024, 7:36pm

This topic was automatically closed 2 days after the last reply. New replies are no longer allowed.

Topic		Replies	Views
Testing streamlit using Mock Using Streamlit session-state	3	2081	March 22, 2024
Setup and tear down test resource in streamlit with pytest Using Streamlit streamlit-cloud , testing , build-with-streamlit	2	767	August 7, 2024
Pytest + session state = no key storage Using Streamlit session-state	2	1407	August 22, 2023
St.session-state has too much dependencies as a data class Random session-state , discussion	5	271	September 8, 2024
Using MagicMock and AppTest Using Streamlit testing	3	1038	August 10, 2024

Mocking out a postgres connection in session state with pytest/monkeypatch

Related topics