Mocking dcim.devices.get and virtualization.devices.get

[rgruyters]

I'm trying to create a mocking for a class function which uses both dcim.devices.get and virtualization.devices.get, but I can't get it to work.

here is an example of my class:

import pynetbox

class Netbox(object):
    def __init__(self, netbox_url: str, netbox_token: str | None = None):
        self.netbox = pynetbox.api(netbox_url, netbox_token)

    def get_device(**kwargs) -> dict:
        device = self.netbox.dcim.devices.get(**kwargs)

       # check if we get a dcim device in return data
        if device is None:
            device = self.netbox.virtualization.devices.get(**kwargs)

        # check if we get a virtual device in the return data
        if device is None:
            return {}

        # return device data as dict
        return dict(device)

Basically, i want to test when self.netbox.dcim.devices.get is None and test if self.netbox.virtualization.devices.get is None.

I have tried the following, but that just simulates the respond data from the request.

from project.netbox import Netbox
from unittest.mock import patch, Mock

nb = Netbox("http://netbox.url", "1234567890abcdef")

@patch("requests.sessions.Session.get")
def test_get_device(mocker):
  mock_response = Mock()
  mock_response.return_value = {"id": 1, "name": "device1"}
  mocker.return_value = mock_response

  response = nb.get_device(name="device1")

  assert response == {"id": 1, "name": "device1"}

Anybody got an idea how to do this?

[markkuleinio]

from project.netbox import Netbox


nb = Netbox("http://netbox.url", "1234567890abcdef")


def test_get_device_device(mocker):
    mock = mocker.Mock()
    mock.dcim.devices.get.return_value = {"id": 1, "name": "device1"}
    mock.virtualization.devices.get.return_value = {"id": 2, "name": "device1"}
    mocker.patch.object(nb, "netbox", mock)
    response = nb.get_device(name="device1")
    assert response == {"id": 1, "name": "device1"}


def test_get_device_vm(mocker):
    mock = mocker.Mock()
    mock.dcim.devices.get.return_value = None
    mock.virtualization.devices.get.return_value = {"id": 2, "name": "device1"}
    mocker.patch.object(nb, "netbox", mock)
    response = nb.get_device(name="device1")
    assert response == {"id": 2, "name": "device1"}

[rgruyters]

Thanks @markkuleinio