diff --git a/tests/test_integration.py b/tests/test_integration.py
index 89701b6..891eb2b 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -17,10 +17,11 @@
 import sys
 import unittest
 import random
+import math
 
-import neal
 import dimod
 import numpy as np
+from dwave.samplers import SimulatedAnnealingSampler
 
 import demo
 import demo_numpy
@@ -58,8 +59,8 @@ def test_num_new_cs(self):
 
         bqm = demo.build_bqm(potential_new_cs_nodes, num_poi, pois, num_cs, charging_stations, num_new_cs)
 
-        sampler = neal.SimulatedAnnealingSampler()
-        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes)
+        sampler = SimulatedAnnealingSampler()
+        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes, num_reads=10, seed=42)
 
         self.assertEqual(num_new_cs, len(new_charging_nodes))
 
@@ -77,8 +78,8 @@ def test_close_to_pois(self):
         bqm = demo.build_bqm(potential_new_cs_nodes, num_poi, pois, num_cs, charging_stations, num_new_cs)
 
         # random.seed(1)
-        sampler = neal.SimulatedAnnealingSampler()
-        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes, seed=42)
+        sampler = SimulatedAnnealingSampler()
+        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes, num_reads=10, seed=42)
 
         new_cs_x = new_charging_nodes[0][0]
         new_cs_y = new_charging_nodes[0][1]
@@ -97,13 +98,10 @@ def test_solution_quality(self):
         bqm = demo.build_bqm(potential_new_cs_nodes, num_poi, pois, num_cs, charging_stations, num_new_cs)
 
         # random.seed(1)
-        sampler = neal.SimulatedAnnealingSampler()
-        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes, seed=1)
+        sampler = SimulatedAnnealingSampler()
+        new_charging_nodes = demo.run_bqm_and_collect_solutions(bqm, sampler, potential_new_cs_nodes, num_reads=10, seed=42)
 
-        new_cs_dist = 0
-        for i in range(num_new_cs):
-            for j in range(i+1, num_new_cs):
-                new_cs_dist += abs(new_charging_nodes[i][0]-new_charging_nodes[j][0])+abs(new_charging_nodes[i][1]-new_charging_nodes[j][1])
+        new_cs_dist = math.sqrt(demo.distance(new_charging_nodes[0], new_charging_nodes[1]))
 
         self.assertGreater(new_cs_dist, 10)