diff --git a/setup.py b/setup.py
index da0d92f..1db0b09 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 import os
 from setuptools import setup
 
-version = '0.1.0'
+version = '0.1.1'
 
 classifiers = [
     "Development Status :: 3 - Alpha",
diff --git a/tests/test.py b/tests/test.py
index f886de6..2d0c1df 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -5,10 +5,10 @@
 from django.test import TestCase
 from django.core.exceptions import ValidationError
 from django_postgresql_dag.exceptions import NodeNotReachableException, GraphModelsCannotBeParsedException, IncorrectUsageException
-from django_postgresql_dag.transformations import _filter_order, edges_from_nodes_queryset, nodes_from_edges_queryset, nx_from_queryset, model_to_dict
+from django_postgresql_dag.transformations import _ordered_filter, edges_from_nodes_queryset, nodes_from_edges_queryset, nx_from_queryset, model_to_dict
+from django_postgresql_dag.query_builders import AncestorQuery, DescendantQuery, UpwardPathQuery, DownwardPathQuery, ConnectedGraphQuery
 
-# from .dag_output import expected_dag_output
-from .models import NetworkNode, NetworkEdge
+from .models import NetworkNode, NetworkEdge, NodeSet, EdgeSet
 
 logging.basicConfig(level=logging.DEBUG)