import unittest

import matplotlib

from pymatgen.util.plotting import periodic_table_heatmap, van_arkel_triangle
from pymatgen.util.testing import PymatgenTest


class FuncTestCase(PymatgenTest):
    def test_plot_periodic_heatmap(self):
        random_data = {
            "Te": 0.11083818874391202,
            "Au": 0.7575629917425387,
            "Th": 1.2475885304040335,
            "Ni": -2.0354391922547705,
        }
        plt = periodic_table_heatmap(random_data, cmap="plasma")
        plt = periodic_table_heatmap(random_data)
        plt = periodic_table_heatmap(random_data, max_row=7)
        plt = periodic_table_heatmap(random_data, max_row=10)
        plt = periodic_table_heatmap(random_data, cbar_label_size=18)
        plt = periodic_table_heatmap(random_data, cmap_range=[0, 1])
        plt = periodic_table_heatmap(random_data, cbar_label="Hello World")
        plt = periodic_table_heatmap(random_data, blank_color="white")
        plt = periodic_table_heatmap(random_data, value_format="%.4f")
        plt = periodic_table_heatmap(random_data, edge_color="black")
        plt = periodic_table_heatmap(random_data, value_fontsize=12)
        plt = periodic_table_heatmap(random_data, symbol_fontsize=18)

    def test_van_arkel_triangle(self):
        random_list = [("Fe", "C"), ("Ni", "F")]
        plt = van_arkel_triangle(random_list)
        plt = van_arkel_triangle(random_list, annotate=True)


if __name__ == "__main__":
    unittest.main()

@@ -14,27 +14,8 @@
 st%0A%0A
-import matplotlib%0A%0A
 from
@@ -401,35 +401,33 @@
       %7D%0A        
-plt
+_
  = periodic_tabl
@@ -464,35 +464,33 @@
 lasma%22)%0A        
-plt
+_
  = periodic_tabl
@@ -512,35 +512,33 @@
 m_data)%0A        
-plt
+_
  = periodic_tabl
@@ -571,35 +571,33 @@
 _row=7)%0A        
-plt
+_
  = periodic_tabl
@@ -631,35 +631,33 @@
 row=10)%0A        
-plt
+_
  = periodic_tabl
@@ -699,35 +699,33 @@
 ize=18)%0A        
-plt
+_
  = periodic_tabl
@@ -766,35 +766,33 @@
 %5B0, 1%5D)%0A        
-plt
+_
  = periodic_tabl
@@ -840,35 +840,33 @@
 World%22)%0A        
-plt
+_
  = periodic_tabl
@@ -909,35 +909,33 @@
 white%22)%0A        
-plt
+_
  = periodic_tabl
@@ -978,35 +978,33 @@
 %22%25.4f%22)%0A        
-plt
+_
  = periodic_tabl
@@ -1046,35 +1046,33 @@
 black%22)%0A        
-plt
+_
  = periodic_tabl
@@ -1121,19 +1121,17 @@
         
-plt
+_
  = perio
@@ -1180,16 +1180,89 @@
 size=18)
+%0A        _ = periodic_table_heatmap(random_data, readable_fontcolor=True)
 %0A%0A    de
@@ -1343,35 +1343,33 @@
 , %22F%22)%5D%0A        
-plt
+_
  = van_arkel_tri
@@ -1399,11 +1399,9 @@
     
-plt
+_
  = v

add test