diff --git a/.travis.yml b/.travis.yml
index d8657a8f..bb7d875e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,13 +8,13 @@ python:
 install:
   - pip install Sphinx sphinx_rtd_theme codecov packaging
   - "python -c $'import os, packaging.version as version\\nv = version.parse(os.environ.get(\"TRAVIS_TAG\", \"1.0\")).public\\nwith open(\"VERSION\", \"w\") as f: f.write(v)'"
-  - python setup.py install
+  - pip install -e .[test]
   - cd docs
   - make clean html
   - cd ..
 
 script:
-  - python setup.py nosetests --with-coverage --cover-package=graphkit
+  - pytest -v --cov=graphkit
 
 deploy:
   provider: pypi
diff --git a/setup.py b/setup.py
index bd7883f4..c8d231a5 100644
--- a/setup.py
+++ b/setup.py
@@ -28,11 +28,15 @@
      author_email='huyng@yahoo-inc.com',
      url='http://github.com/yahoo/graphkit',
      packages=['graphkit'],
-     install_requires=['networkx'],
+     install_requires=[
+          "networkx; python_version >= '3.5'",
+          "networkx == 2.2; python_version < '3.5'",
+     ],
      extras_require={
-          'plot': ['pydot', 'matplotlib']
+          'plot': ['pydot', 'matplotlib'],
+          'test': ['pydot', 'matplotlib', 'pytest', "pytest-cov"],
      },
-     tests_require=['numpy'],
+     tests_require=['pytest', "pytest-cov"],
      license='Apache-2.0',
      keywords=['graph', 'computation graph', 'DAG', 'directed acyclical graph'],
      classifiers=[
diff --git a/test/test_graphkit.py b/test/test_graphkit.py
index bd97b317..7db2e973 100644
--- a/test/test_graphkit.py
+++ b/test/test_graphkit.py
@@ -6,7 +6,8 @@
 
 from pprint import pprint
 from operator import add
-from numpy.testing import assert_raises
+
+import pytest
 
 import graphkit.network as network
 import graphkit.modifiers as modifiers
@@ -180,9 +181,10 @@ def test_pruning_raises_for_bad_output():
 
     # Request two outputs we can compute and one we can't compute.  Assert
     # that this raises a ValueError.
-    assert_raises(ValueError, net, {'a': 1, 'b': 2, 'c': 3, 'd': 4},
-        outputs=['sum1', 'sum3', 'sum4'])
-
+    with pytest.raises(ValueError) as exinfo:
+        net({'a': 1, 'b': 2, 'c': 3, 'd': 4},
+            outputs=['sum1', 'sum3', 'sum4'])
+    assert exinfo.match('sum4')
 
 def test_optional():
     # Test that optional() needs work as expected.