Original notebook of custom_enums_and_tests.md

custom_enums_and_tests.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Motivation\n",
    "\n",
    "Assume you want to share across your application a smart list of constants.\n",
    "For example, a list of colors supported by your application.\n",
    "One way probably could be a `list` in some module along with some replated logic implemented.\n",
    "In this post, you will learn how I tackled this issue.\n",
    "\n",
    "## Enums in Python\n",
    "\n",
    "There are a lot of resources about Enums in Python and probably the best place to start are the [docs](https://docs.python.org/3/library/enum.html).\n",
    "Let us start with a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "class Color(Enum):\n",
    "    RED = 1\n",
    "    GREEN = 2\n",
    "    BLUE = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is great; you can check the value of `RED`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Color.RED: 1>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Color.RED"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall, that the objective is to come up with a solution enabling a little more logic.\n",
    "In particular, you want to have two functionalities:\n",
    "\n",
    "1. Nice printout of the supported colors\n",
    "1. Checking whether a given color is supported\n",
    "\n",
    "## Customizing `Color`\n",
    "\n",
    "Here is the definition of the customized class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import EnumMeta\n",
    "\n",
    "\n",
    "class MyEnumMeta(EnumMeta):\n",
    "    def __str__(cls):\n",
    "        lines = [f\"Members of `{cls.__name__}` are:\"]\n",
    "        for member in cls:\n",
    "            lines.append(f\"- {member}\")\n",
    "        return '\\n'.join(lines)\n",
    "\n",
    "    def _contains(self, member):\n",
    "        return member in self._member_map_ \\\n",
    "            or member in set(\n",
    "                map(lambda x: x.value, self._member_map_.values()))\n",
    "\n",
    "    def is_valid(self, member):\n",
    "        if self._contains(member):\n",
    "            return True\n",
    "        else:\n",
    "            return False\n",
    "\n",
    "\n",
    "class Color(Enum, metaclass=MyEnumMeta):\n",
    "    RED = 1\n",
    "    GREEN = 2\n",
    "    BLUE = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, your `Color` can do much more:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Members of `Color` are:\n",
      "- Color.RED\n",
      "- Color.GREEN\n",
      "- Color.BLUE\n"
     ]
    }
   ],
   "source": [
    "print(Color)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or, alternatively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, False)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Color.is_valid('RED'), Color.is_valid('YELLOW')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I am sure you can take it from here, improve and modify the implementation as per your needs.\n",
    "But, before doing so, you might want to see how the custom `Color` can be used, and, more importantly, tested.\n",
    "\n",
    "## Using the custom `Color`\n",
    "\n",
    "Here is an example of how we can use the custom `Color`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def foo(type):\n",
    "    if Color.is_valid(type):\n",
    "        print(f\"Wonderful! we will use the wonderful {type} color\")\n",
    "        return \"valid\"\n",
    "    else:\n",
    "        print(f\"{type} is not supported color. The supported colors are:\\n{str(Color)}\")\n",
    "        return \"invalid\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here are some examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wonderful! we will use the wonderful RED color\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'valid'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo('RED')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wonderful! we will use the wonderful BLUE color\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'valid'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo('BLUE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "YELLOW is not supported color. The supported colors are:\n",
      "Members of `Color` are:\n",
      "- Color.RED\n",
      "- Color.GREEN\n",
      "- Color.BLUE\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'invalid'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "foo('YELLOW')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing `foo`\n",
    "\n",
    "Next, you want to test `foo`.\n",
    "But, as you want it to be a proper unit test, you don't want to refer to the values of `Color`.\n",
    "That is, you want to patch `Color` and use some value defined within the scope of the test.\n",
    "Here is the way to do that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import unittest\n",
    "from unittest.mock import patch\n",
    "\n",
    "\n",
    "class DummyColor(Enum, metaclass=MyEnumMeta):\n",
    "    TEST_TYPE1 = 1\n",
    "    TEST_TYPE2 = 10\n",
    "\n",
    "\n",
    "class MyTest(unittest.TestCase):\n",
    "\n",
    "    def test_a(self):\n",
    "        # Without any patch, you are depending on the \n",
    "        # values of `Color`\n",
    "        assert foo('RED') == 'valid'\n",
    "        assert foo('YELLOW') == 'invalid'\n",
    "\n",
    "    @patch('__main__.Color', DummyColor)  # Make sure you patch the right `Color`\n",
    "                                          # it should be within the class you want to test\n",
    "    def test_b(self):\n",
    "        assert foo('TEST_TYPE1') == 'valid'\n",
    "        assert foo('TEST_TYPE10') == 'invalid'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".."
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wonderful! we will use the wonderful RED color\n",
      "YELLOW is not supported color. The supported colors are:\n",
      "Members of `Color` are:\n",
      "- Color.RED\n",
      "- Color.GREEN\n",
      "- Color.BLUE\n",
      "Wonderful! we will use the wonderful TEST_TYPE1 color\n",
      "TEST_TYPE10 is not supported color. The supported colors are:\n",
      "Members of `DummyColor` are:\n",
      "- DummyColor.TEST_TYPE1\n",
      "- DummyColor.TEST_TYPE2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "----------------------------------------------------------------------\n",
      "Ran 2 tests in 0.002s\n",
      "\n",
      "OK\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    unittest.main(argv=['first-arg-is-ignored'], exit=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "I hope you found it useful.\n",
    "As a bonus, here is more or less the same example, given in scripts.\n",
    "You can save the three of theme somewhere and run `pytest` to see how everything works together.\n",
    "\n",
    "### `special_enum.py`\n",
    "\n",
    "```python\n",
    "# Implementation of the customized enum\n",
    "from enum import Enum, EnumMeta\n",
    "\n",
    "\n",
    "class MyEnumMeta(EnumMeta):\n",
    "    def __str__(cls):\n",
    "        lines = [f\"Members of `{cls.__name__}` are:\"]\n",
    "        for member in cls:\n",
    "            lines.append(f\"- {member}\")\n",
    "        return '\\n'.join(lines)\n",
    "\n",
    "    def _contains(self, member):\n",
    "        return member in self._member_map_ \\\n",
    "            or member in set(\n",
    "                map(lambda x: x.value, self._member_map_.values()))\n",
    "\n",
    "    def is_valid(self, member):\n",
    "        if self._contains(member):\n",
    "            return True\n",
    "        else:\n",
    "            return False\n",
    "\n",
    "\n",
    "class SpecialEnum(Enum, metaclass=MyEnumMeta):\n",
    "    TYPE1 = 1\n",
    "    TYPE2 = 10\n",
    "```\n",
    "\n",
    "### `foo.py`\n",
    "\n",
    "```python\n",
    "# Your amazing logic!\n",
    "from special_enum import SpecialEnum\n",
    "\n",
    "\n",
    "def bar(type):\n",
    "    print(SpecialEnum)  # Just for debugging\n",
    "    if SpecialEnum.is_valid(type):\n",
    "        return \"valid\"\n",
    "    else:\n",
    "        return \"invalid\"\n",
    "```\n",
    "\n",
    "### `test_foo.py`\n",
    "\n",
    "```python\n",
    "# The testing script\n",
    "import unittest\n",
    "from unittest.mock import patch\n",
    "import foo\n",
    "from special_enum import MyEnumMeta\n",
    "from enum import Enum\n",
    "\n",
    "\n",
    "class SpecialEnumTest(Enum, metaclass=MyEnumMeta):\n",
    "    TEST_TYPE1 = 1\n",
    "    TEST_TYPE2 = 10\n",
    "\n",
    "\n",
    "class MyTest(unittest.TestCase):\n",
    "\n",
    "    def test_a(self):\n",
    "        # Passes as it is ussing the original SpecialEnum\n",
    "        assert foo.bar('TYPE1') == 'valid'\n",
    "        assert foo.bar('TYPE10') == 'invalid'\n",
    "\n",
    "    @patch('foo.SpecialEnum', SpecialEnumTest)\n",
    "    def test_b(self):\n",
    "        assert foo.bar('TEST_TYPE1') == 'valid'\n",
    "        assert foo.bar('TEST_TYPE10') == 'invalid'\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}