test_userfactory.py

import unittest
from userfactory import UserFactory


class TestUserFactory(unittest.TestCase):
    def setUp(self):
        self.__factory = UserFactory()

    def test_factory_creates_staff(self):
        staff = self.__factory.factory("staff")
        role = staff.get_role()
        self.assertEqual(role, "staff")

    def test_factory_creates_user(self):
        staff = self.__factory.factory("user")
        role = staff.get_role()
        self.assertEqual(role, "user")