diff --git a/.github/workflows/deploy.yaml b/.github/workflows/deploy.yaml
index f155390..79dc15d 100644
--- a/.github/workflows/deploy.yaml
+++ b/.github/workflows/deploy.yaml
@@ -1,5 +1,4 @@
-name: Publish Python 🐍 distributions 📦 to PyPI
-
+name: Deploy
 on:
   push:
     tags:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 6f06f25..370887f 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -1,15 +1,13 @@
-name: Python package
+name: CI
 
 on: [ pull_request ]
 
 jobs:
-  build:
-
+  test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: [ "3.7", "3.8", "3.9", "3.10" ]
-
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
diff --git a/README.md b/README.md
index 348e408..50ffa2b 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,7 @@
 # pytest-sqlalchemy-mock
 [![PyPI version](https://badge.fury.io/py/pytest-sqlalchemy-mock.svg)](https://badge.fury.io/py/pytest-sqlalchemy-mock)
+![CI](https://github.com/resulyrt93/pytest-sqlalchemy-mock/actions/workflows/tests.yml/badge.svg?branch=dev)
+<a href="https://github.com/psf/black"><img alt="Code style: black" src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
 
 This plugin provides pytest fixtures to create an in-memory DB instance on tests and dump your raw test data.
 
diff --git a/pytest_sqlalchemy_mock/__init__.py b/pytest_sqlalchemy_mock/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pytest_sqlalchemy_mock.py b/pytest_sqlalchemy_mock/base.py
similarity index 57%
rename from pytest_sqlalchemy_mock.py
rename to pytest_sqlalchemy_mock/base.py
index 67ec730..d684e23 100644
--- a/pytest_sqlalchemy_mock.py
+++ b/pytest_sqlalchemy_mock/base.py
@@ -1,40 +1,15 @@
-from typing import TYPE_CHECKING, List, Tuple
+from typing import TYPE_CHECKING
 
 import pytest
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
+from .model_mocker import ModelMocker
+
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
 
-class ModelMocker:
-    # TODO make type annotations
-    _base = None
-    _mock_config: List[Tuple[str, List]] = None
-
-    def __init__(self, session, base, mock_config):
-        self._session: Session = session
-        self._base = base
-        self._mock_config = mock_config
-
-    def get_model_class_with_table_name(self, table_name: str):
-        for mapper in self._base.registry.mappers:
-            cls = mapper.class_
-            if cls.__tablename__ == table_name:
-                return cls
-
-    def create_all(self):
-        for model_config in self._mock_config:
-            table_name, data = model_config
-            model_class = self.get_model_class_with_table_name(table_name)
-            if model_class:
-                for datum in data:
-                    instance = model_class(**datum)
-                    self._session.add(instance)
-        self._session.commit()
-
-
 @pytest.fixture(scope="session")
 def connection_url():
     return "sqlite:///:memory:"
diff --git a/pytest_sqlalchemy_mock/model_mocker.py b/pytest_sqlalchemy_mock/model_mocker.py
new file mode 100644
index 0000000..f414037
--- /dev/null
+++ b/pytest_sqlalchemy_mock/model_mocker.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import List, Tuple
+
+    from sqlalchemy.orm import Session
+
+    MOCK_CONFIG_TYPE = List[Tuple[str, List]]
+
+
+class ModelMocker:
+    _base = None
+    _mock_config: MOCK_CONFIG_TYPE = None
+
+    def __init__(self, session: Session, base, mock_config: MOCK_CONFIG_TYPE):
+        self._session: Session = session
+        self._base = base
+        self._mock_config = mock_config
+
+    def get_model_class_with_table_name(self, table_name: str):
+        for mapper in self._base.registry.mappers:
+            cls = mapper.class_
+            if cls.__tablename__ == table_name:
+                return cls
+
+    def create_all(self):
+        for model_config in self._mock_config:
+            table_name, data = model_config
+            model_class = self.get_model_class_with_table_name(table_name)
+            if model_class:
+                for datum in data:
+                    instance = model_class(**datum)
+                    self._session.add(instance)
+        self._session.commit()
diff --git a/tests/conftest.py b/tests/conftest.py
index 6816aef..fcbc4f9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,5 @@
-import pytest
-
-from pytest_sqlalchemy_mock import *
+# TODO use pytester for testing
+from pytest_sqlalchemy_mock.base import *
 from tests.data import MockData
 from tests.db import Base