cleanup fixture

README.md

$ python -m pytest -rP test1.py
================================== test session starts ==================================
platform darwin -- Python 3.8.13, pytest-7.2.0, pluggy-1.0.0
rootdir: /Users/xavier/src/streamsets/stf-4.x, configfile: pytest.ini
collected 2 items

test1.py::test_with_s3_buckets PASSED                                             [ 50%]
test1.py::test_with_bad_resource SKIPPED (unconditional skip)                     [100%]

======================================== PASSES =========================================
_________________________________ test_with_s3_buckets __________________________________
--------------------------------- Captured stdout call ----------------------------------
Allocating temp_file temp_file1
Allocating temp_file temp_file2
Allocating temp_file temp_file3
code using temp_file temp_file3
Freeing temp_file temp_file3
Allocating bucket first_bucket
code using bucket first_bucket
Freeing bucket first_bucket
Allocating bucket second_bucket
Allocating bucket third_bucket
Done!
------------------------------- Captured stdout teardown --------------------------------
Freeing bucket third_bucket
Raisin while cleaning poor_resource foo
Exception in __Exit__!
Freeing bucket second_bucket
Freeing temp_file temp_file2
Freeing temp_file temp_file1
============================= 1 passed, 1 skipped in 0.01s ==============================

conftest.py

import contextlib

import pytest

def writeOutput(output):
    print(output)
    # with open('test1.log', 'a') as f:
    #    f.write(f'{output}\n')

@pytest.fixture
def sch():
    """A fixture we depend on"""
    yield


class CleanUp(object):
    """Helper class for the cleanup fixture"""
    def __init__(self, sch, exit_stack):
        self.on_exit = exit_stack

    def _track(self, resource):
      if not hasattr(resource, '__enter__') or not hasattr(resource, '__exit__'):
        raise TypeError(f"Track expects a generator function")
      return self.on_exit.enter_context(NoExceptionsOnCleanup(resource))

    def auto_cleanup(self, resource_factory):
      def wrapper(*args, **kwargs):
        return self._track(resource_factory(*args, **kwargs))
      return wrapper

      
class NoExceptionsOnCleanup(contextlib.AbstractContextManager):
    """Wrap a context manager so that exceptions are suppressed"""
    def __init__(self, cm):
        self.cm = cm

    def __enter__(self):
        return self.cm.__enter__()

    def __exit__(self, exc_type, exc_value, traceback):
        try:
          report = self.cm.__exit__(exc_type, exc_value, traceback)
          if exc_type and not report: 
            writeOutput('Context would have not suppressed the exception')
        except:
          writeOutput('Exception in __Exit__!')
        return True


@pytest.fixture
def cleanup(sch):
    """
    Fixture that provides and object to help in test cleanup.
     `cleanup.auto_cleanup` wraps a function creating a context manager so
         that the context it creates are cleaned up when exiting the test
    ~~`cleanup.track` adds to the to-clean list a resource that is created using
         a context manager, and returns that resourc~~
    ~~`cleanup.on_exit` provides an ExitStack to add functions to be executed
      when the test completes~~
      See https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
    """
    with contextlib.ExitStack() as exit_stack:
      yield CleanUp(sch, exit_stack)

test1.py

import contextlib

import pytest

def writeOutput(output):
    print(output)
    # with open('test1.log', 'a') as f:
    #    f.write(f'{output}\n')

@contextlib.contextmanager
def s3_bucket(sch, name):
  writeOutput(f'Allocating bucket {name}')
  yield f'bucket {name}'
  writeOutput(f'Freeing bucket {name}')

@contextlib.contextmanager
def temp_file(sch, name):
  writeOutput(f'Allocating temp_file {name}')
  yield f'temp_file {name}'
  writeOutput(f'Freeing temp_file {name}')

@contextlib.contextmanager
def poor_resource(name):
  yield f'poor resource {name}'
  writeOutput(f'Raisin while cleaning poor_resource {name}')
  raise ValueError("Poor resource")


def test_with_s3_buckets(sch, cleanup):

  create_temp_file = cleanup.auto_cleanup(temp_file)
  create_s3_bucket = cleanup.auto_cleanup(s3_bucket)
  create_broken_resource = cleanup.auto_cleanup(poor_resource)


  temp_file1 = create_temp_file(sch, 'temp_file1')
  temp_file2 = create_temp_file(sch, 'temp_file2')
  with temp_file(sch, 'temp_file3') as temp_file3:
    writeOutput(f'code using {temp_file3}')

  with s3_bucket(sch, 'first_bucket') as first_bucket:
    writeOutput(f'code using {first_bucket}')

  second_bucket = create_s3_bucket(sch, 'second_bucket')
  create_broken_resource('foo')
  third_bucket = create_s3_bucket(sch, 'third_bucket')

  writeOutput(f'Done!')

@pytest.mark.skip
def test_with_bad_resource(cleanup):
    def _plain_function():
        return "plain_function"

    cleanup.track(_plain_function)