From ffa73a47eef034722a4d6c89993144b6ee48b473 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Sat, 25 Dec 2021 09:27:02 -0500
Subject: [PATCH] Add DiscreteConditional unit test for markdown printing

---
 .../gtsam/tests/test_DiscreteConditional.py   | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 python/gtsam/tests/test_DiscreteConditional.py

diff --git a/python/gtsam/tests/test_DiscreteConditional.py b/python/gtsam/tests/test_DiscreteConditional.py
new file mode 100644
index 000000000..5e24dc40b
--- /dev/null
+++ b/python/gtsam/tests/test_DiscreteConditional.py
@@ -0,0 +1,54 @@
+"""
+GTSAM Copyright 2010-2021, Georgia Tech Research Corporation,
+Atlanta, Georgia 30332-0415
+All Rights Reserved
+
+See LICENSE for the license information
+
+Unit tests for Discrete Conditionals.
+Author: Varun Agrawal
+"""
+
+# pylint: disable=no-name-in-module, invalid-name
+
+import unittest
+
+from gtsam import DiscreteConditional, DiscreteKeys
+from gtsam.utils.test_case import GtsamTestCase
+
+
+class TestDiscreteConditional(GtsamTestCase):
+    """Tests for Discrete Conditionals."""
+    def test_markdown(self):
+        """Test whether the _repr_markdown_ method."""
+
+        A = (2, 2)
+        B = (1, 2)
+        C = (0, 3)
+        parents = DiscreteKeys()
+        parents.push_back(B)
+        parents.push_back(C)
+
+        conditional = DiscreteConditional(A, parents,
+                                          "0/1 1/3  1/1 3/1  0/1 1/0")
+        expected = \
+            " $P(A|B,C)$:\n" \
+            "|B|C|0|1|\n" \
+            "|:-:|:-:|:-:|:-:|\n" \
+            "|0|0|0|1|\n" \
+            "|0|1|0.25|0.75|\n" \
+            "|0|2|0.5|0.5|\n" \
+            "|1|0|0.75|0.25|\n" \
+            "|1|1|0|1|\n" \
+            "|1|2|1|0|\n"
+
+        def formatter(x: int):
+            names = ["C", "B", "A"]
+            return names[x]
+
+        actual = conditional._repr_markdown_(formatter)
+        self.assertEqual(actual, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()