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

Using Streamlit
,
1

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?

2

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.

3

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.

4

Same thing. That code does noting but defining Test_Burger .

5

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.

6

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"]
7

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

8

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.

9

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).

10

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.

1 Like
11

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.

12

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.

13

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!

14

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