import os
from importlib import import_module

from .testdriver import TestDriver


class Driver(TestDriver):
    drivers = {}

    def __new__(cls, type, *args, **kwargs):
        if type not in cls.drivers:
            try:
                mod = import_module('onitu.drivers.{}.tests.driver'.
                                    format(type))
            except ImportError:
                raise KeyError("No such driver {}".format(repr(type)))
            cls.drivers[type] = mod.Driver
        return cls.drivers[type](*args, **kwargs)


class LocalStorageDriver(TestDriver):
    def __new__(cls, *args, **kwargs):
        return Driver('local_storage', *args, **kwargs)


class TargetDriver(Driver):
    def __new__(cls, *args, **kwargs):
        type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
        return Driver(type, *args, **kwargs)

@@ -7,43 +7,28 @@
  os%0A
-from 
 import
-lib import import_module
+ pkg_resources
 %0A%0Afr
@@ -127,19 +127,19 @@
 __(cls, 
-typ
+nam
 e, *args
@@ -163,176 +163,538 @@
     
-if type not in cls.drivers:%0A            try:%0A                mod = import_module('onitu.drivers.%7B%7D.tests.driver'.%0A                                    format(type))%0A    
+entry_points = pkg_resources.iter_entry_points('onitu.tests')%0A        tests_modules = %7Be.name: e for e in entry_points%7D%0A%0A        if name not in tests_modules:%0A            raise ImportError(%0A                %22Cannot import tests for driver %7B%7D%22.format(name)%0A            )%0A%0A        try:%0A            tests = tests_modules%5Bname%5D.load()%0A        except ImportError as e:%0A            raise ImportError(%0A                %22Error importing tests for driver %7B%7D: %7B%7D%22.format(name, e)%0A            )%0A%0A        try:%0A            driver = tests.Driver%0A
     
@@ -733,71 +733,139 @@
     
-    
 raise 
-KeyError(%22No such driver %7B%7D%22.format(repr(type)))%0A    
+ImportError(%0A                %22Tests for driver %7B%7D don't expose a%22%0A                %22Driver class%22.format(name)%0A            )%0A%0A
     
@@ -884,21 +884,17 @@
 ers%5B
-typ
+nam
 e%5D = 
-mod.D
+d
 rive
@@ -914,25 +914,14 @@
 urn 
-cls.
 driver
-s%5Btype%5D
 (*ar

Load tests helpers using entry_points