File: test_distribution_plots.py

package info (click to toggle)
python-cogent 1.9-14
  • links: PTS, VCS
  • area: main
  • in suites: buster
  • size: 19,752 kB
  • sloc: python: 137,485; makefile: 149; sh: 64
file content (521 lines) | stat: -rwxr-xr-x 24,823 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
#!/usr/bin/env python
__author__ = "Jai Ram Rideout"
__copyright__ = "Copyright 2007-2016, The Cogent Project"
__credits__ = ["Jai Ram Rideout"]
__license__ = "GPL"
__version__ = "1.9"
__maintainer__ = "Jai Ram Rideout"
__email__ = "jai.rideout@gmail.com"
__status__ = "Production"

"""Tests public and private functions in the distribution_plots module."""
from warnings import filterwarnings

filterwarnings("ignore", "More than 20 figures")

from matplotlib import use
use('Agg', warn=False)

from StringIO import StringIO
import sys
import matplotlib.colors as colors
from matplotlib.pyplot import boxplot
from numpy import array
from cogent.draw.distribution_plots import _validate_input,\
    _get_distribution_markers, _validate_x_values, _create_plot,\
    _calc_data_point_locations, _set_axes_options, generate_box_plots,\
    generate_comparative_plots, _calc_data_point_ticks, _plot_bar_data,\
    _plot_scatter_data, _plot_box_data, _color_box_plot,\
    _create_legend, _set_figure_size
from cogent.util.unit_test import TestCase, main

class DistributionPlotsTests(TestCase):
    """Tests of the distribution_plots module."""
    def setUp(self):
        """Create some data to be used in the tests."""
        # Test null data list.
        self.Null = None

        # Test empty data list.
        self.Empty = []

        # Test nested empty data list.
        self.EmptyNested = [[]]

        # Test nested empty data list (for bar/scatter plots).
        self.EmptyDeeplyNested = [[[]]]

        # Test invalid number of samples in data list (for bar/scatter plots).
        self.InvalidNumSamples = [[[1, 2, 3, 4, 5]],
                                  [[4, 5, 6, 7, 8], [2, 3, 2]],
                                  [[4, 7, 10, 33, 32, 6, 7, 8]]]

        # Test valid data with one sample (for bar/scatter plots).
        self.ValidSingleSampleData = [[[1, 2, 3, 4, 5]],
                                      [[4, 5, 6, 7, 8]],
                                      [[4, 7, 10, 33, 32, 6, 7, 8]]]

        # Test valid data with three samples and four data points
        # (for bar/scatter plots).
        self.ValidTypicalData = [[[1.0, 2, 3.5, 5], [2, 3, 5, 6], [2, 3, 8]],
                                 [[4, 7, 8], [8, 9, 10, 11], [9.0, 4, 1, 1]],
                                 [[4, 33, 32, 6, 8], [5, 4, 8, 13], [1, 1, 2]],
                                 [[2, 2, 2, 2], [3, 9, 8], [2, 1, 6, 7, 4, 5]]]

        # Test typical data to be plotted by the boxplot function.
        self.ValidTypicalBoxData = [[3.4, 10, 11.67, 12.0, 2, 2, 99.99],
                                    [2.3, 4, 5, 88, 9, 10, 11, 1, 0, 3, -8],
                                    [2, 9, 7, 5, 6]]

    def test_validate_input_null(self):
        """_validate_input() should raise a ValueError if null data is passed
        to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.Null, None, None, None)

    def test_validate_input_empty(self):
        """_validate_input() should raise a ValueError if empty data is passed
        to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.Empty, None, None, None)

    def test_validate_input_empty_nested(self):
        """_validate_input() should raise a ValueError if empty nested data is
        passed to it."""
        self.assertRaises(ValueError, _validate_input,
                          self.EmptyNested, None, None, None)

    def test_validate_input_empty_deeply_nested(self):
        """_validate_input() should pass for deeply nested empty data."""
        num_points, num_samples = _validate_input(self.EmptyDeeplyNested,
                                                  None, None, None)
        self.assertEqual(num_points, 1)
        self.assertEqual(num_samples, 1)

    def test_validate_input_invalid_num_samples(self):
        """_validate_input() should raise a ValueError if an inconsistent
        number of samples in included in the data."""
        self.assertRaises(ValueError, _validate_input,
                          self.InvalidNumSamples, None, None, None)

    def test_validate_x_values_invalid_x_values(self):
        """_validate_x_values() should raise a ValueError on an invalid number
        of x_values."""
        self.assertRaises(ValueError, _validate_x_values,
                          [1, 2, 3, 4], ["T0", "T1", "T2"],
                          len(self.ValidSingleSampleData))

    def test_validate_x_values_invalid_x_tick_labels(self):
        """_validate_x_values() should raise a ValueError on an invalid number
        of x_tick_labels."""
        self.assertRaises(ValueError, _validate_x_values,
                          None, ["T0"], len(self.ValidSingleSampleData))

    def test_validate_x_values_nonnumber_x_values(self):
        """_validate_x_values() should raise a ValueError on x_values that
        aren't numbers."""
        self.assertRaises(ValueError, _validate_x_values,
                ["foo", 2, 3], None, len(self.ValidSingleSampleData))

    def test_validate_x_values_valid_x_values(self):
        """_validate_x_values() should not throw an exception."""
        _validate_x_values([1, 2.0, 3], None, 3)

    def test_validate_input_invalid_data_point_names(self):
        """_validate_input() should raise a ValueError on data_point_names that
        are an invalid length."""
        self.assertRaises(ValueError, _validate_input,
                self.ValidSingleSampleData, None, ["T0", "T1"], None)

    def test_validate_input_invalid_sample_names(self):
        """_validate_input() should raise a ValueError on sample_names that are
        an invalid length."""
        self.assertRaises(ValueError, _validate_input,
                self.ValidSingleSampleData, None, None, ["Men", "Women"])

    def test_validate_input_all_valid_input(self):
        """_validate_input() should return valid information about the data
        without throwing an exception."""
        self.assertEqual(_validate_input(self.ValidTypicalData, [1, 3, 4, 8],
                                         ["T0", "T1", "T2", "T3"],
                                         ["Infants", "Children", "Teens"]),
                                         (4, 3))

    def test_get_distribution_markers_null_marker_list(self):
        """_get_distribution_markers() should return a list of predefined
        matplotlib markers."""
        self.assertEqual(_get_distribution_markers('colors', None, 5),
                ['b', 'g', 'r', 'c', 'm'])

    def test_get_distribution_markers_empty_marker_list(self):
        """_get_distribution_markers() should return a list of predefined
        matplotlib markers."""
        self.assertEqual(_get_distribution_markers('colors', None, 4),
                ['b', 'g', 'r', 'c'])

    def test_get_distribution_markers_insufficient_markers(self):
        """_get_distribution_markers() should return a wrapped list of
        predefined markers."""
        # Save stdout and replace it with something that will capture the print
        # statement. Note: this code was taken from here:
        # http://stackoverflow.com/questions/4219717/how-to-assert-output-
        #     with-nosetest-unittest-in-python/4220278#4220278
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            self.assertEqual(_get_distribution_markers('colors', None, 10),
                ['b', 'g', 'r', 'c', 'm', 'y', 'w', 'b', 'g', 'r'])
            self.assertEqual(_get_distribution_markers('symbols',
                ['^', '>', '<'], 5), ['^', '>', '<', '^', '>'])
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.\nThere are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_get_distribution_markers_bad_marker_type(self):
        """_get_distribution_markers() should raise a ValueError."""
        self.assertRaises(ValueError, _get_distribution_markers, 'shapes', [],
                3)

    def test_get_distribution_markers_zero_markers(self):
        """_get_distribution_markers() should return an empty list."""
        self.assertEqual(_get_distribution_markers('symbols', None, 0), [])
        self.assertEqual(_get_distribution_markers('symbols', ['^'], 0), [])

    def test_create_plot(self):
        """_create_plot() should return a tuple containing a Figure and
        Axes."""
        fig, ax = _create_plot()
        self.assertEqual(fig.__class__.__name__, "Figure")
        self.assertEqual(ax.__class__.__name__, "AxesSubplot")

    def test_plot_bar_data(self):
        """_plot_bar_data() should return a list of Rectangle objects."""
        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [1, 2, 3], 'red', 0.5, 3.75, 1.5, 'stdv')
        self.assertEqual(result[0].__class__.__name__, "Rectangle")
        self.assertEqual(len(result), 1)
        self.assertFloatEqual(result[0].get_width(), 0.5)
        self.assertFloatEqual(result[0].get_facecolor(), (1.0, 0.0, 0.0, 1.0))
        self.assertFloatEqual(result[0].get_height(), 2.0)

        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [1, 2, 3], 'red', 0.5, 3.75, 1.5, 'sem')
        self.assertEqual(result[0].__class__.__name__, "Rectangle")
        self.assertEqual(len(result), 1)
        self.assertFloatEqual(result[0].get_width(), 0.5)
        self.assertFloatEqual(result[0].get_facecolor(), (1.0, 0.0, 0.0, 1.0))
        self.assertFloatEqual(result[0].get_height(), 2.0)

    def test_plot_bar_data_bad_error_bar_type(self):
        """_plot_bar_data() should raise an exception on bad error bar type."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _plot_bar_data, ax, [1, 2, 3], 'red',
                0.5, 3.75, 1.5, 'var')

    def test_plot_bar_data_empty(self):
        """_plot_bar_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [], 'red', 0.5, 3.75, 1.5, 'stdv')
        self.assertTrue(result is None)

        fig, ax = _create_plot()
        result = _plot_bar_data(ax, [], 'red', 0.5, 3.75, 1.5, 'sem')
        self.assertTrue(result is None)

    def test_plot_scatter_data_empty(self):
        """_plot_scatter_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_scatter_data(ax, [], '^', 0.77, 1, 1.5, 'stdv')
        self.assertTrue(result is None)

    def test_plot_box_data(self):
        """_plot_box_data() should return a dictionary for Line2D's."""
        fig, ax = _create_plot()
        result = _plot_box_data(ax, [0, 0, 7, 8, -3, 44], 'blue', 0.33, 55,
                1.5, 'stdv')
        self.assertEqual(type(result), dict)
        # make sure the values are not empty
        for key in ['boxes', 'medians', 'whiskers', 'fliers', 'caps']:
            val = result[key]
            self.assertTrue(len(val), 0)

    def test_plot_box_data_empty(self):
        """_plot_box_data() should not error when given empty list of data,
        but should not plot anything."""
        fig, ax = _create_plot()
        result = _plot_box_data(ax, [], 'blue', 0.33, 55, 1.5, 'stdv')
        # the original test was poor as it was testing a property of matplotlib
        # I'm simply exercising this here. These capabilities will be dropped
        # in a future version.
        for key in ['boxes', 'medians', 'whiskers', 'fliers', 'caps']:
            val = result[key]

    def test_calc_data_point_locations_invalid_widths(self):
        """_calc_data_point_locations() should raise a ValueError
        exception when it encounters bad widths."""
        self.assertRaises(ValueError, _calc_data_point_locations, [1, 2, 3],
                3, 2, -2, 0.5)
        self.assertRaises(ValueError, _calc_data_point_locations, [1, 2, 3],
                3, 2, 2, -0.5)

    def test_calc_data_point_locations_default_spacing(self):
        """_calc_data_point_locations() should return an array containing
        the x-axis locations for each data point, evenly spaced from 1..n."""
        locs = _calc_data_point_locations(None, 4, 2, 0.25, 0.5)
        self.assertEqual(locs, array([1.0, 2.0, 3.0, 4.0]))

    def test_calc_data_point_locations_custom_spacing(self):
        """_calc_data_point_locations() should return an array containing
        the x-axis locations for each data point, spaced according to a custom
        spacing scheme."""
        locs = _calc_data_point_locations([3, 4, 10, 12], 4, 2, 0.25, 0.75)
        self.assertEqual(locs, array([3.75, 5.0, 12.5, 15.0]))

    def test_calc_data_point_ticks(self):
        """_calc_data_point_ticks() should return an array containing the
        x-axis locations for each data point tick."""
        ticks = _calc_data_point_ticks(array([1, 5, 9, 11]), 1, 0.5, False)
        self.assertFloatEqual(ticks, array([1.25, 5.25, 9.25, 11.25]))

        ticks = _calc_data_point_ticks(array([0]), 3, 0.5, False)
        self.assertFloatEqual(ticks, array([0.75]))

    def test_set_axes_options(self):
        """_set_axes_options() should set the labels on the axes and not raise
        any exceptions."""
        fig, ax = _create_plot()
        _set_axes_options(ax, "Plot Title", "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1"])
        self.assertEqual(ax.get_title(), "Plot Title")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(ax.get_xticklabels()[0].get_text(), "T0")
        self.assertEqual(ax.get_xticklabels()[1].get_text(), "T1")

    def test_set_axes_options_ylim(self):
        """_set_axes_options() should set the y-axis limits."""
        fig, ax = _create_plot()
        _set_axes_options(ax, "Plot Title", "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1", "T2"], y_min=0, y_max=1)
        self.assertEqual(ax.get_title(), "Plot Title")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(ax.get_xticklabels()[0].get_text(), "T0")
        self.assertEqual(ax.get_xticklabels()[1].get_text(), "T1")
        self.assertFloatEqual(ax.get_ylim(), [0, 1])

    def test_set_axes_options_bad_ylim(self):
        """_set_axes_options() should raise an exception when given non-numeric
        y limits."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _set_axes_options, ax, "Plot Title",
                          "x-axis label", "y-axis label",
                          x_tick_labels=["T0", "T1", "T2"], y_min='car',
                          y_max=30)

    def test_create_legend(self):
        """_create_box_plot_legend() should create a legend on valid input."""
        fig, ax = _create_plot()
        _create_legend(ax, ['b', 'r'], ['dist1', 'dist2'], 'colors')
        self.assertEqual(len(ax.get_legend().get_texts()), 2)

        fig, ax = _create_plot()
        _create_legend(ax, ['^', '<', '>'], ['dist1', 'dist2', 'dist3'],
                'symbols')
        self.assertEqual(len(ax.get_legend().get_texts()), 3)

    def test_create_legend_invalid_input(self):
        """Test raises error on bad input."""
        fig, ax = _create_plot()
        self.assertRaises(ValueError, _create_legend, ax,
                ['^', '<', '>'], ['dist1', 'dist2'], 'symbols')
        self.assertRaises(ValueError, _create_legend, ax, ['^', '<', '>'],
                ['dist1', 'dist2', 'dist3'], 'foo')

    def test_generate_box_plots(self):
        """generate_box_plots() should return a valid Figure object."""
        fig = generate_box_plots(self.ValidTypicalBoxData, [1, 4, 10],
                                 ["Data 1", "Data 2", "Data 3"], "Test",
                                 "x-axis label", "y-axis label")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 3)
        self.assertFloatEqual(ax.get_xticks(), [1, 4, 10])

    def test_generate_comparative_plots_bar(self):
        """generate_comparative_plots() should return a valid barchart Figure
        object."""
        fig = generate_comparative_plots('bar', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'r', 'g'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.3, 7.4, 17.6, 19.3])

    def test_generate_comparative_plots_insufficient_colors(self):
        """generate_comparative_plots() should work even when there aren't
        enough colors. We should capture a print statement that warns the
        users."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            generate_comparative_plots('bar', self.ValidTypicalData,
                    [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                    ["Infants", "Children", "Teens"], ['b', 'r'],
                    "x-axis label", "y-axis label", "Test")
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_generate_comparative_plots_scatter(self):
        """generate_comparative_plots() should return a valid scatterplot
        Figure object."""
        fig = generate_comparative_plots('scatter', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['^', '>', '<'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.1, 7.2, 17.4, 19.1])

    def test_generate_comparative_plots_insufficient_symbols(self):
        """generate_comparative_plots() should work even when there aren't
        enough symbols. We should capture a print statement that warns the
        users."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out
            generate_comparative_plots('scatter', self.ValidTypicalData,
                    [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                    ["Infants", "Children", "Teens"], ['^'],
                    "x-axis label", "y-axis label", "Test")
            output = out.getvalue().strip()
            self.assertEqual(output, "There are not enough markers to "
                    "uniquely represent each distribution in your dataset. "
                    "You may want to provide a list of markers that is at "
                    "least as large as the number of distributions in your "
                    "dataset.")
        finally:
            sys.stdout = saved_stdout

    def test_generate_comparative_plots_empty_marker_list(self):
        """generate_comparative_plots() should use the predefined list of
        markers if an empty list is provided by the user."""
        generate_comparative_plots('scatter', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], [],
                "x-axis label", "y-axis label", "Test")

    def test_generate_comparative_plots_box(self):
        """generate_comparative_plots() should return a valid boxplot Figure
        object."""
        fig = generate_comparative_plots('box', self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'g', 'y'],
                "x-axis label", "y-axis label", "Test")
        ax = fig.get_axes()[0]
        self.assertEqual(ax.get_title(), "Test")
        self.assertEqual(ax.get_xlabel(), "x-axis label")
        self.assertEqual(ax.get_ylabel(), "y-axis label")
        self.assertEqual(len(ax.get_xticklabels()), 4)
        self.assertFloatEqual(ax.get_xticks(), [2.1, 7.2, 17.4, 19.1])

    def test_generate_comparative_plots_error(self):
        """generate_comparative_plots() should raise a ValueError for an
        invalid plot type."""
        self.assertRaises(ValueError, generate_comparative_plots, 'pie',
                self.ValidTypicalData,
                [1, 4, 10, 11], ["T0", "T1", "T2", "T3"],
                ["Infants", "Children", "Teens"], ['b', 'g', 'y'],
                "x-axis label", "y-axis label", "Test")

    def test_color_box_plot(self):
        """_color_box_plot() should not throw an exception when passed the
        proper input."""
        fig, ax = _create_plot()
        box_plot = boxplot(self.ValidTypicalBoxData)
        _color_box_plot(ax, box_plot, 'blue')

    def test_set_figure_size(self):
        """Test setting a valid figure size."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        _set_figure_size(fig, 3, 4)
        self.assertFloatEqual(fig.get_size_inches(), (3, 4))

    def test_set_figure_size_defaults(self):
        """Test setting a figure size using matplotlib defaults."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        orig_fig_size = fig.get_size_inches()
        _set_figure_size(fig)
        self.assertFloatEqual(fig.get_size_inches(), orig_fig_size)

    def test_set_figure_size_invalid(self):
        """Test setting a figure size using invalid dimensions."""
        fig, ax = _create_plot()
        _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                          x_tick_labels=['foofoofoo', 'barbarbar'],
                          x_tick_labels_orientation='vertical')
        orig_fig_size = fig.get_size_inches()
        _set_figure_size(fig, -1, 0)
        self.assertFloatEqual(fig.get_size_inches(), orig_fig_size)

    def test_set_figure_size_long_labels(self):
        """Test setting a figure size that has really long labels."""
        saved_stdout = sys.stdout
        try:
            out = StringIO()
            sys.stdout = out

            fig, ax = _create_plot()
            _set_axes_options(ax, 'foo', 'x_foo', 'y_foo',
                              x_tick_labels=['foofoofooooooooooooooooooooooooo'
                              'ooooooooooooooooooooooooooooooooooooooooooooooo'
                              'ooooooooooooooooooooo', 'barbarbar'],
                              x_tick_labels_orientation='vertical')
            _set_figure_size(fig, 3, 3)
            self.assertFloatEqual(fig.get_size_inches(), (3, 3))
            output = out.getvalue().strip()
#            self.assertEqual(output,
#            "Warning: could not automatically resize plot to make room for "
#            "axes labels and plot title. This can happen if the labels or "
#            "title are extremely long and the plot size is too small. Your "
#            "plot may have its labels and/or title cut-off. To fix this, "
#            "try increasing the plot's size (in inches) and try again.")
        finally:
            sys.stdout = saved_stdout


if __name__ == '__main__':
    main()