From ff8f7802c829bdf683440ff28636bc041474c2e5 Mon Sep 17 00:00:00 2001
From: Victoria Goliber <vgoliber@dwavesys.com>
Date: Mon, 24 May 2021 14:19:48 +0200
Subject: [PATCH] Address feedback

---
 nurse_scheduling.py            | 13 ++++++++++---
 requirements.txt               |  3 +++
 tests/test_nurse_scheduling.py |  8 +-------
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/nurse_scheduling.py b/nurse_scheduling.py
index b92982d..8dc4555 100644
--- a/nurse_scheduling.py
+++ b/nurse_scheduling.py
@@ -24,8 +24,16 @@
 from dimod import BinaryQuadraticModel
 from collections import defaultdict
 from copy import deepcopy
-import matplotlib.pyplot as plt
-from matplotlib.patches import Rectangle
+import matplotlib
+
+try:
+    import matplotlib.pyplot as plt
+    from matplotlib.patches import Rectangle
+except ImportError:
+    matplotlib.use("agg")
+    import matplotlib.pyplot as plt
+    from matplotlib.patches import Rectangle
+
 
 # Overall model variables: problem size
 # binary variable q_nd is the assignment of nurse n to day d
@@ -172,7 +180,6 @@ def get_nurse_and_day(index):
 
 # Get the results
 smpl = results.first.sample
-# energy = results.first.energy
 
 # Graphics
 print("\nBuilding schedule and checking constraints...\n")
diff --git a/requirements.txt b/requirements.txt
index c244693..01bd522 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,4 @@
 dwave-ocean-sdk>=3.3.0
+
+matplotlib==3.3.4; python_version>'3.5'
+matplotlib==3.0.3; python_version=='3.5'
diff --git a/tests/test_nurse_scheduling.py b/tests/test_nurse_scheduling.py
index 9df08d3..dfb0e5d 100644
--- a/tests/test_nurse_scheduling.py
+++ b/tests/test_nurse_scheduling.py
@@ -19,13 +19,7 @@
 example_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 class IntegrationTests(unittest.TestCase):
-    @unittest.skipIf(os.getenv('SKIP_INT_TESTS'), "Skipping integration test.")
-    def test_smoke(self):
-        """run nurse_scheduling.py and check that nothing crashes"""
-
-        demo_file = os.path.join(example_dir, 'nurse_scheduling.py')
-        subprocess.check_output([sys.executable, demo_file])
-    
+    @unittest.skipIf(os.getenv('SKIP_INT_TESTS'), "Skipping integration test.")    
     def test_schedule(self):
         file_path = os.path.join(example_dir, 'nurse_scheduling.py')