From dd4cd41c317f1649a691e1b236ebbb51eb26daa8 Mon Sep 17 00:00:00 2001
From: Mark Piper <mark.piper@colorado.edu>
Date: Tue, 5 Sep 2023 16:12:09 -0700
Subject: [PATCH] Add tests for Diffusion class

---
 .../ivy-diffusion/tests/test_diffusion.py     | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 lessons/python/ivy-diffusion/tests/test_diffusion.py

diff --git a/lessons/python/ivy-diffusion/tests/test_diffusion.py b/lessons/python/ivy-diffusion/tests/test_diffusion.py
new file mode 100644
index 0000000..456450f
--- /dev/null
+++ b/lessons/python/ivy-diffusion/tests/test_diffusion.py
@@ -0,0 +1,31 @@
+"""Tests for the diffusion module."""
+import numpy as np
+
+from ivy_diffusion.diffusion import Diffusion
+
+SHAPE = 10
+SPACING = 1.0
+DIFFUSIVITY = 1.0
+ZMAX = 5.0
+
+
+def test_defaults():
+    m = Diffusion()
+    assert type(m) == Diffusion
+    assert m.shape == SHAPE
+    assert m.spacing == SPACING
+    assert m.diffusivity == DIFFUSIVITY
+    assert m.time == 0.0
+    assert m.time_step > 0.0
+    assert type(m.concentration) == np.ndarray
+    assert len(m.concentration) == SHAPE
+
+
+def test_update():
+    m = Diffusion()
+    m.concentration = np.zeros(SHAPE)
+    m.concentration[SHAPE//2] = ZMAX
+    m.update()
+    assert m.time > 0.0
+    assert m.concentration.max() < ZMAX
+    assert m.concentration.sum() == ZMAX