Test correct model selection in a DRF View

I have a view with a get endpoint which determines which model to use by looking if the query contains a "bla" query param:

class FooView(APIView):
    def get(self, request):

        if "bla" in self.request.query_params:
            model = ModelA
        else:
            model = ModelB
        
        # ...

Now I want to test this via pytest. The approach below seems to work. Is the use of method_calls correct here? Nevertheless, it looks a bit clumsy.

@pytest.mark.django_db
def test_get_foo_correct_model_selection(client):
    with mock.patch('path.to.ModelA.objects') as model_a_manager, mock.patch('path.to.ModelB.objects') as model_b_manager:

        response = client.get(reverse("foo-list") + "?bla=true")

        assert len(model_a_manager.method_calls) >= 1
        assert len(model_b_manager.method_calls) == 0


    # and now without bla param
    with mock.patch('path.to.ModelA.objects') as model_a_manager, mock.patch('path.to.ModelB.objects') as model_b_manager:
        response = client.get(reverse("foo-list"))

        assert len(model_a_manager.method_calls) == 0
        assert len(model_b_manager.method_calls) >= 1

Is there a more elegant way?

Usually I would mock method that accepts your queryset or model instance as param and then you can use:

# checks if method was called with correct param.
mocked_method.assert_called_with(correct_model) 

In your example you can use

model_a_manager.assert_called()
model_b_manager.assert_not_called()
Back to Top