Remove/overwrite import

Question

I'm trying to set up a grading script for an intro CS class using unittest. Essentially, students submit a python file student.py, which has some number of functions in it that are generally interdependent (meaning that func3() may use func1() in it's computations).

I'm writing unit tests for each method by comparing the output of student.func1 to the output of correct.func1, a method that's known to be the correct implementation (from a file correct.py).

As an example, say that func2 uses func1 in it's computation. So, either by default or upon student.func1 failing some test, I want to overwrite student.func1 with correct.func1, so student.func2 uses a known correct implementation (and therefore isn't just wrong by default). How could I go about doing that? It seems like the setUp() and tearDown() are similar to what I want, but I don't know how to "unimport" modules in python, and haven't found any resources about it so far.

I'm interested in both the case where student.py contains classes, and func1,func2 are methods of a specific class, and when func1 and func2 are just defined generically within student.py.

Answer 1

The easiest way would be to import student into your module, then catch an AssertionError if a test fails, and replace the bad code inside the student module with your own good code:

import student

import unittest

def safe_f1():
    print("Safe f1")
    return 1

class TestSomething(unittest.TestCase):

    def test_f1(self):
        try:
            self.assertEqual(student.func1(), 1)
        except AssertionError:
            student.func1 = safe_f1
            raise


    def test_f2(self):
        self.assertEqual(student.func2(), 2)

Here's a dummy student.py that fails/works:

def func1():
    print("Bad f1")
    return 2

def func2():
    return func1() + 1
    return 2

When I run this, I get:

$ python -m unittest test.py
Bad f1
FSafe f1
.
======================================================================
FAIL: test_f1 (test.TestSomething)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/austin/Code/so/test.py", line 13, in test_f1
    self.assertEqual(student.func1(), 1)
AssertionError: 2 != 1

----------------------------------------------------------------------
Ran 2 tests in 0.001s

FAILED (failures=1)