Mock pystalkd Beanstalkd connection for unit testing in Django - Python

fakes.py

class FakeBeanstalkdJob():
    body = None
    def __init__(self, body=''):
        if not isinstance(body, str):
            raise Exception('Job body must be a str instance')
        self.body = body

    def delete(self):
        pass

    def bury(self):
        pass


class FakeBeanstalkd():
    queue = dict()
    tube = None
    
    def __init__(self, ip=None, port=None):
        return self

    def use(tube):
        FakeBeanstalkd.tube = tube
        if tube not in FakeBeanstalkd.queue:
            FakeBeanstalkd.queue[tube] = list()
        return

    def put(data):
        if FakeBeanstalkd.tube is None:
            raise Exception("No tube selected via `use`!")

        job = FakeBeanstalkdJob(data)
        FakeBeanstalkd.queue[FakeBeanstalkd.tube].append(job)
        return

    def watch(tube):
        FakeBeanstalkd.tube = tube
        if tube not in FakeBeanstalkd.queue:
            FakeBeanstalkd.queue[tube] = list()
        return

    def reserve(index=0):
        if FakeBeanstalkd.tube is None:
            raise Exception("No tube selected via `watch`!")

        number_of_jobs = len(FakeBeanstalkd.queue[FakeBeanstalkd.tube])
        
        if number_of_jobs == 0:
            return None

        job = FakeBeanstalkd.queue[FakeBeanstalkd.tube].pop(0)
        return job

    def kick(max):
        pass

test.py

from django.test import TestCase
from fakes import FakeBeanstalkd

class MyAppTestCase(TestCase):
    @mock.patch('pystalkd.Beanstalkd.Connection', FakeBeanstalkd)
    def test_put(self):
        conn = pystalkd.Beanstalkd.Connection('127.0.0.1', 27017)
        conn.use('my-queue1')
        conn.put('some-data-blabla')
        conn.put('some-data-bliblibli')
        conn.watch('my-queue1')
        job = conn.reserve(0)
        print(job.body)