How to unit test mocking SQLAlchemy engine in Python for an update query

Question

I have a function called update_worker_data. This simply updates a table in PostgreSQL, however if the row does not exist in the db, it then inserts it. This is checked by getting the rowcount after the query is executed:

def update_worker_data(db: engine, data: List[dict]) -> int:
    """Update the data for the modifiable table.

    Args:
    ----------
    * db(engine): db engine
    * date(List[dict]): data to either be updated or inserted  

    Returns:
    ----------
    * row_count(int): amount of rows updated and inserted 

    """
    update_query = """
    UPDATE worker_data
    SET "first_col" = %(f_col)s, "second_col" = %(s_col)s, "third_col" = %(t_col)s, "fourth_col" = %(fo_col)s
    WHERE alt_id = %(a_id)s
    """
    insert_query = """
    INSERT INTO worker_data("first_col", "second_col", "third_col", "fourth_col", "alt_id")
    VALUES (%(f_col)s, %(s_col)s, %(t_col)s, %(fo_col)s, %(a_id)s);
    """
    row_count = 0
    for d in data:
        params = {
            "f_col": d["first_col"],
            "s_col": d["second_col"],
            "t_col": d["third_col"],
            "fo_col": d["fourth_col"],
            "a_id": d["alt_id"]
        }
        with db.connect() as conn:
            has_updated = conn.execute(update_query, params).rowcount
            has_inserted = 0
            if has_updated == 0:
                has_inserted = conn.execute(insert_query, params).rowcount
            row_count += (has_inserted + has_updated)
    return row_count

My unit test is currently like so:

@patch("api.db.engine")
def test_update_c_data(mock_engine, update_data_dict):
    cursor_mock = mock_engine.connect.return_value.__enter__.return_value
    cursor_mock.execute.return_value.rowcount = 1
    actual_row_count = update_substance_store(mock_engine, update_data_dict)
    assert actual_row_count == 4

When I run this unit test, it asserts true, however this will only ever execute the first query (update_query) and will never enter the has_update == 0 block due to returning value of .rowcount equalling 1. So according to the mock docs I should use side_effect. I then create a list of values for the side_effect and use parameterize. My desired operation would then to be run first the update_query and then the insert_query.

I then changed my unit test to use side_effect so it becomes like so:

@pytest.mark.parameterize("expected_row_count", [
    ([1, 0]),
    ([0, 1]),
])
@patch("api.db.engine")
def test_update_c_data(mock_engine, expected_row_count, update_data_dict):
    cursor_mock = mock_engine.connect.return_value.__enter__.return_value
    cursor_mock.execute.return_value.rowcount.side_effect = expected_row_count 
    actual_row_count = update_substance_store(mock_engine, update_data_dict)
    assert actual_row_count == 5

However when I run this test, it just sets the value to a mock object. Is side_effect the correct way to handle this test case?