from pillar.tests import PillarTestServer, AbstractPillarTest
from attract_server import AttractExtension


class AttractTestServer(PillarTestServer):
    def __init__(self, *args, **kwargs):
        super(AttractTestServer, self).__init__(*args, **kwargs)

        self.load_extension(AttractExtension(), '/attract')


class AbstractAttractTest(AbstractPillarTest):
    pillar_server_class = AttractTestServer