Mocking internals of a Python Script

How to test against a Python script when you need to change how __main__ behaves

This will be split into two parts. The first contains the contents of some script.py file, which is the base template script we want to use:

Code

#| language: python
def test_function(): 
    return 2

def main():
    result = test_function()
    with open("someFile.txt", "w") as f:
        f.write(str(result))
        
if __name__ == "__main__":
    main()

Code + Explanation

#| language: python
def test_function(): 
    return 2

def main():
    result = test_function()
    with open("someFile.txt", "w") as f:
        f.write(str(result))
        
if __name__ == "__main__":
    main()

---

def test_function(): 
    return 2

This is the function whose behavior we want to override during our mock

---

def main():
    result = test_function()
    with open("someFile.txt", "w") as f:
        f.write(str(result))

The main function is what we will call when running the python script. This will write a string representation of test_function() to a file.

---

if __name__ == "__main__":
    main()

This is a blocker to ensure that if anyone imports or calls this python script that it will be ran explicitly.

File Structure

It should be assumed that for the next part the structure of the code files are as such:

• base_repository
  • example
    • script.py
  • tests
    • test_script.py

Code

#| language: python
import os
import sys
import unittest
from unittest import mock
SRC_DIRS = [
    os.path.join(
        os.path.dirname(__file__), "example"
    )
]
sys.path.extend(SRC_DIRS)
if SRC_DIRS is not None:
    import script
def new_function():
    return 0
@mock.patch("script.test_function", new_function)
class ExampleTester(unittest.TestCase):
    def test_example(self):
        script.main()
        with open("someFile.txt", "r") as f:
            lines = f.read()
        self.assertEquals(lines, "0")

Code + Explanation

#| language: python
import os
import sys
import unittest
from unittest import mock
SRC_DIRS = [
    os.path.join(
        os.path.dirname(__file__), "example"
    )
]
sys.path.extend(SRC_DIRS)
if SRC_DIRS is not None:
    import script
def new_function():
    return 0
@mock.patch("script.test_function", new_function)
class ExampleTester(unittest.TestCase):
    def test_example(self):
        script.main()
        with open("someFile.txt", "r") as f:
            lines = f.read()
        self.assertEquals(lines, "0")

---

import os
import sys
import unittest
from unittest import mock

These are the imports we use

---

SRC_DIRS = [
    os.path.join(
        os.path.dirname(__file__), "example"
    )
]

This is a list of directories that have our script source code relative to the current file. In this case the example directory.

---

sys.path.extend(SRC_DIRS)

We add in our new SRC_DIRS to the sys.path which allows them to be imported through an import statement such as import script

---

if SRC_DIRS is not None:
    import script

If the file exists (this makes it modular) go ahead and import it

---

def new_function():
    return 0

This is the new function we will use to replace the test_function in our python script

---

@mock.patch("script.test_function", new_function)

This uses unittest.mock to mokey-patch and override the original test_function in the existing module with the new one we just defined. Calling script.test_function() will call new_function() as a result

---

class ExampleTester(unittest.TestCase):
    def test_example(self):
        script.main()

Calls the main function in our tester, but uses our new_function() when called

---

        with open("someFile.txt", "r") as f:
            lines = f.read()
        self.assertEquals(lines, "0")

Tests that the file which was written to has the properly mocked version of it, or 0