How to Unit Test a function that use NATS KV

[Oceanestars]

Hi, I'm looking to Mock the nats server and be able to interact it in my unit test.(Using pytest)

I have a NATS wrapper so I can easily grab the key from the nats bucket

    async def get_key(self, key: str, bucket: str = None) -> bytes:
        """Use KV proxy to get the value for this key from nats KV storage"""
        if bucket is None:
            bucket = f"{self.get_target_source_type()}_bucket"
        if bucket not in self.kv_proxies:
            self.kv_proxies[bucket] = KVProxy(bucket)
        return await self.kv_proxies[bucket].get(key)

The function I want to test is

    async def get_data(self, key str) -> Dict:
        try:
            data_binary = await self.get_key(
                key, "data-25"
            )
            return json.loads(data_binary.decode())
        except KeyNotFoundError:
            return None

I'm trying to use https://github.com/nats-io/nats.py/blob/main/tests/utils.py

When looking for some NATS unit test example, google AI came up with this:

import unittest
from unittest.mock import patch, AsyncMock
import asyncio
from nats.aio.client import Client as NATS

class MyNATSClient:
    async def connect_to_nats(self):
        self.nc = NATS()
        await self.nc.connect("nats://localhost:4222")

    async def publish_message(self, subject, message):
        await self.nc.publish(subject, message.encode())

    async def close_connection(self):
        await self.nc.close()

class TestMyNATSClient(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
      self.nats_client = MyNATSClient()

    async def asyncTearDown(self):
        if hasattr(self.nats_client, 'nc') and self.nats_client.nc.is_connected:
            await self.nats_client.close_connection()
    
    @patch('nats.aio.client.Client.connect', new_callable=AsyncMock)
    async def test_connect_to_nats(self, mock_connect):
        await self.nats_client.connect_to_nats()
        mock_connect.assert_awaited_once_with("nats://localhost:4222")

It ran but couldn't figure out how to connect it to the get_data function. I thought setting the connection as a python fixture. I wasn't successful either.

Any ideas?

Thanks for any help!