# Copyright 2019 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import csv
import importlib.resources
import io
import json
import os
import pathlib
import re
import statistics
import subprocess
import sys
import tarfile
import tempfile
import unittest
import unittest.mock
from typing import Any, Callable, Dict, List, Tuple, Union

import perfcompare

# The following allows perfcompare_test.py to be run either:
#
# 1. Using Fuchsia's GN plumbing for Python code, where test data must
#    be imported through a special module using importlib.  This case
#    allows the test to be run on Fuchsia Infra.
# 2. By running perfcompare_test.py directly through
#    fuchsia-vendored-python or python3.  This case allows the golden file
#    to be updated, by passing "--generate".
try:
    import perfcompare_test_data
except ModuleNotFoundError:
    RESOURCES_DIR = pathlib.Path(__file__).parent
else:
    RESOURCES_DIR = importlib.resources.files(perfcompare_test_data)

Behavior = Union[
    str,
    Callable[[List[str]], subprocess.CompletedProcess],
    Exception,
]


# Test case helper class for creating temporary directories that will
# be cleaned up when the test finishes.
class TempDirTestCase(unittest.TestCase):
    def MakeTempDir(self):
        return self.enterContext(
            tempfile.TemporaryDirectory(
                prefix="tmp_unittest_%s_" % self.__class__.__name__
            )
        )


def WriteJsonFile(filename, json_data):
    with open(filename, "w") as fh:
        json.dump(json_data, fh)


def ReadGoldenFile(file_obj):
    with file_obj.open("r") as fh:
        data = fh.read()
    matches = list(re.finditer("\n\n### (.*)\n", data, re.M))
    starts = [m.end() for m in matches]
    ends = [m.start() for m in matches[1:]] + [len(data)]
    for m, start, end in zip(matches, starts, ends):
        yield m.group(1), data[start:end]


# Helper for checking against test expectations in a golden file.
# This provides an implementation of AssertCaseEq() that compares
# results against the golden file.
class GoldenDataInput(object):
    def __init__(self, filename):
        self._cases = dict(ReadGoldenFile(filename))

    def AssertCaseEq(self, name, actual):
        expected = self._cases[name]
        if expected != actual:
            raise AssertionError('"%s" != "%s"' % (actual, expected))


# This provides an implementation of AssertCaseEq() that updates the
# golden file with new expectations generated by the tests.
class GoldenDataOutput(object):
    def __init__(self):
        self._cases = {}

    def AssertCaseEq(self, name, actual):
        assert name not in self._cases, name
        self._cases[name] = actual

    def WriteFile(self, file_obj):
        with file_obj.open("w") as fh:
            for name, data in sorted(self._cases.items()):
                fh.write("\n\n### %s\n%s" % (name, data))


GOLDEN_FILE = RESOURCES_DIR.joinpath("perfcompare_test_output.txt")
GOLDEN = GoldenDataInput(GOLDEN_FILE)


def TestMain():
    global GOLDEN
    if "--generate" in sys.argv:
        sys.argv.pop(sys.argv.index("--generate"))
        GOLDEN = GoldenDataOutput()
        try:
            unittest.main()
        finally:
            GOLDEN.WriteFile(GOLDEN_FILE)
    else:
        unittest.main()


# Test data from a normal distribution, generated using the following code:
# ', '.join('%.4f' % random.gauss(0, 1) for _ in xrange(100))
TEST_VALUES = [
    0.4171,
    2.1056,
    -0.0223,
    -1.6592,
    0.4766,
    -0.6405,
    0.3488,
    1.5729,
    2.0654,
    -0.1324,
    -0.8648,
    -0.2793,
    -0.7966,
    0.2851,
    -0.9374,
    -2.0275,
    0.8222,
    -0.2396,
    -0.6982,
    0.9067,
    0.9416,
    -2.2870,
    -0.1868,
    1.0700,
    -1.2531,
    0.8455,
    1.4755,
    0.2979,
    0.3441,
    0.6694,
    -0.1808,
    -0.9038,
    0.8267,
    -0.4320,
    -0.7166,
    0.3757,
    -0.5135,
    -0.9497,
    2.0372,
    -0.3364,
    0.3879,
    -0.2970,
    1.3872,
    0.6538,
    1.0674,
    1.2349,
    -0.6873,
    -0.1807,
    0.6867,
    -0.1150,
    -1.0526,
    -0.6853,
    -0.5858,
    -1.8460,
    1.6041,
    -1.1638,
    0.5459,
    -1.6476,
    -0.8711,
    -0.9001,
    0.0788,
    -0.8170,
    0.2439,
    0.0129,
    -0.8674,
    -1.1076,
    -0.0074,
    -0.6230,
    -0.4761,
    -2.2526,
    0.4906,
    -0.5001,
    -0.2050,
    0.7623,
    -0.5511,
    -0.2837,
    -0.8797,
    -0.5374,
    -1.2910,
    0.9551,
    0.4483,
    -0.6352,
    -0.3334,
    -0.5105,
    0.1073,
    2.9131,
    -0.4941,
    -0.2808,
    -0.2517,
    -1.9961,
    0.9214,
    -0.6325,
    -1.1895,
    0.8118,
    1.5424,
    0.5601,
    -1.0322,
    0.7135,
    -0.2780,
    -0.1128,
]


def GenerateTestData(mean, stddev):
    return [x * stddev + mean for x in TEST_VALUES]


# This is an example of a slow running time value for an initial run of a
# test.  This should be skipped by the software under test.
SLOW_INITIAL_RUN = [1e6]


# Returns the output of the "compare_perf" subcommand when run on the
# given directories.
def ComparePerf(dir_paths, expected_sample_size=100):
    stdout = io.StringIO()
    perfcompare.Main(
        [
            "compare_perf",
            "--expected_sample_size=%d" % expected_sample_size,
        ]
        + dir_paths,
        stdout,
    )
    return stdout.getvalue()


class FormatConfidenceIntervalTest(unittest.TestCase):
    def test_confidence_interval_formatting(self):
        Format = perfcompare.FormatConfidenceInterval

        self.assertEqual(Format(12345.6789, 2222), "12346 +/- 2222")
        self.assertEqual(Format(12345.6789, 0.02222), "12345.679 +/- 0.022")
        self.assertEqual(Format(12345.6789, 0.07777), "12345.679 +/- 0.078")
        self.assertEqual(Format(12345.6789, 0.09911), "12345.679 +/- 0.099")
        # Corner case: rounding 0.09950 to 2 significant figures produces
        # 0.100, which looks like 3 significant figures rather than 2.
        self.assertEqual(Format(12345.6789, 0.09950), "12345.679 +/- 0.100")
        self.assertEqual(Format(12345.6789, 2e-5), "12345.678900 +/- 0.000020")

        # Corner case: the offset is a power of 10.
        self.assertEqual(Format(12345.6789, 0.1), "12345.68 +/- 0.10")
        self.assertEqual(Format(12345.6789, 0.01), "12345.679 +/- 0.010")

        # Corner case: zero offset.
        self.assertEqual(Format(12345.6789, 0), "12345.7 +/- 0")

        # Corner case: negative offset.  This does not make sense for a
        # confidence interval and should not happen, but let's ensure it
        # gets formatted anyway in case that it useful for debugging.
        self.assertEqual(Format(12345.6789, -1), "12345.7 +/- -1")

        # Corner cases: infinity and NaN.
        self.assertEqual(Format(12345.6789, float("inf")), "12345.7 +/- inf")
        self.assertEqual(Format(12345.6789, -float("inf")), "12345.7 +/- -inf")
        self.assertEqual(Format(12345.6789, float("nan")), "12345.7 +/- nan")
        self.assertEqual(Format(float("inf"), 0.1234), "inf +/- 0.12")
        self.assertEqual(Format(-float("inf"), 0.1234), "-inf +/- 0.12")
        self.assertEqual(Format(float("nan"), 0.1234), "nan +/- 0.12")


# Generate some example perf test data, allowing variation at each level of
# the sampling process (per boot, per process, and per iteration within
# each process).  This follows a random effects model.  Returns a list of
# lists of lists of values.
def GenerateData(
    mean=1000,
    stddev_across_boots=0,
    stddev_across_processes=0,
    stddev_across_iters=0,
):
    it = iter(TEST_VALUES)

    def GenerateValues(mean, stddev, count):
        return [next(it) * stddev + mean for _ in range(count)]

    # This reads 4**3 + 4**2 + 4 = 84 values from TEST_VALUES, so it does
    # not exceed the number of values in TEST_VALUES.
    return [
        [
            SLOW_INITIAL_RUN
            + GenerateValues(mean_within_process, stddev_across_iters, 4)
            for mean_within_process in GenerateValues(
                mean_within_boot, stddev_across_processes, 4
            )
        ]
        for mean_within_boot in GenerateValues(mean, stddev_across_boots, 4)
    ]


class StatisticsTest(TempDirTestCase):
    def ResultsDictForValues(self, run_values):
        return {
            "label": "ExampleTest",
            "test_suite": "example_suite",
            "unit": "nanoseconds",
            "values": run_values,
        }

    # Given data in the format returned by GenerateData(), writes this data
    # to a temporary directory.
    def DirOfData(self, data):
        dir_path = self.MakeTempDir()
        os.mkdir(os.path.join(dir_path, "by_boot"))
        for boot_idx, results_for_boot in enumerate(data):
            test_dir = os.path.join(
                dir_path,
                "by_boot",
                "boot%06d" % boot_idx,
                "test-name",
                "subdir",
            )
            os.makedirs(test_dir)
            for process_idx, run_values in enumerate(results_for_boot):
                dest_file = os.path.join(
                    test_dir,
                    "example_process%06d.fuchsiaperf.json" % process_idx,
                )
                WriteJsonFile(
                    dest_file, [self.ResultsDictForValues(run_values)]
                )
        return dir_path

    # Sanity-check that DirOfData() writes data in the correct format by
    # reading back some simple test data.
    def test_readback_of_data(self):
        data = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
        dataset = perfcompare.MultiBootDataset(self.DirOfData(data))
        boot_datasets = list(dataset.GetBootDatasets())
        self.assertEqual(len(boot_datasets), 2)
        self.assertEqual(
            list(boot_datasets[0].GetProcessDatasets()),
            [
                [self.ResultsDictForValues([1, 2])],
                [self.ResultsDictForValues([3, 4])],
            ],
        )
        self.assertEqual(
            list(boot_datasets[1].GetProcessDatasets()),
            [
                [self.ResultsDictForValues([5, 6])],
                [self.ResultsDictForValues([7, 8])],
            ],
        )

    def TarFileOfDir(self, dir_path, write_mode):
        tar_filename = os.path.join(self.MakeTempDir(), "out.tar")
        with tarfile.open(tar_filename, write_mode) as tar:
            for name in os.listdir(dir_path):
                tar.add(os.path.join(dir_path, name), arcname=name)
        return tar_filename

    def test_readback_of_data_from_tar_file(self):
        data = [[[1, 2], [3, 4]]]
        dir_path = self.DirOfData(data)
        self.assertEqual(len(os.listdir(os.path.join(dir_path, "by_boot"))), 1)
        # Test the uncompressed and gzipped cases.
        for write_mode in ("w", "w:gz"):
            tar_filename = self.TarFileOfDir(
                os.path.join(dir_path, "by_boot", "boot000000"), write_mode
            )
            boot_dataset = perfcompare.SingleBootDataset(tar_filename)
            self.assertEqual(
                list(boot_dataset.GetProcessDatasets()),
                [
                    [self.ResultsDictForValues([1, 2])],
                    [self.ResultsDictForValues([3, 4])],
                ],
            )

    def CheckConfidenceInterval(self, data, interval_string):
        dir_path = self.DirOfData(data)
        test_name = "example_suite: ExampleTest"
        stats = perfcompare.StatsFromMultiBootDataset(
            perfcompare.MultiBootDataset(dir_path)
        )[test_name]
        self.assertEqual(stats.FormatConfidenceInterval(), interval_string)

    # Test the CIs produced with variation at different levels of the
    # multi-level sampling process.
    def test_confidence_intervals(self):
        self.CheckConfidenceInterval(GenerateData(), "1000 +/- 0 ns")
        self.CheckConfidenceInterval(
            GenerateData(stddev_across_boots=100), "1021 +/- 452 ns"
        )
        self.CheckConfidenceInterval(
            GenerateData(stddev_across_processes=100), "1012 +/- 151 ns"
        )
        self.CheckConfidenceInterval(
            GenerateData(stddev_across_iters=100), "981 +/- 74 ns"
        )

    # Test the case where just a single value is produced per process run.
    def test_confidence_interval_with_single_value_per_process(self):
        self.CheckConfidenceInterval([[[100]], [[101]]], "100 +/- 32 ns")

    # If the "before" and "after" results have identical confidence
    # intervals, that should be treated as "no difference", including when
    # the CIs are zero-width (as tested here).
    def test_comparing_equal_zero_width_confidence_intervals(self):
        dir_path = self.DirOfData([[[200]], [[200]]])
        output = ComparePerf([dir_path, dir_path], expected_sample_size=2)
        GOLDEN.AssertCaseEq("comparison_no_change_zero_width_ci", output)

    def test_sample_size_check_larger(self):
        dir1_path = self.DirOfData([[[10]]] * 4)
        dir2_path = self.DirOfData([[[10]]] * 5)
        # Check that we get a descriptive error if the actual sample
        # sizes do not match the "--expected_sample_size" argument.
        error = (
            "^The following metrics had an unexpected sample size"
            r" \(expected <=3\):\n"
            r"  example_suite: ExampleTest \(got 4\)\n"
            r"  example_suite: ExampleTest \(got 5\)$"
        )
        with self.assertRaisesRegex(AssertionError, error):
            ComparePerf([dir1_path, dir2_path], expected_sample_size=3)
        # We should get a similar error if we pass only one directory.
        error = (
            "^The following metrics had an unexpected sample size"
            r" \(expected <=3\):\n"
            r"  example_suite: ExampleTest \(got 4\)$"
        )
        with self.assertRaisesRegex(AssertionError, error):
            ComparePerf([dir1_path], expected_sample_size=3)
        # There should be no check for the sample sizes if the
        # "--expected_sample_size" argument is not given.
        stdout = io.StringIO()
        perfcompare.Main(["compare_perf", dir1_path, dir2_path], stdout)

    def test_sample_size_check_smaller(self):
        # We shouldn't get an error if the actual sample sizes are
        # smaller than expected.
        dir1_path = self.DirOfData([[[10]]] * 3)
        dir2_path = self.DirOfData([[[10]]] * 4)
        ComparePerf([dir1_path, dir2_path], expected_sample_size=5)


class PerfCompareTest(TempDirTestCase):
    def AddIgnoredFiles(self, dest_dir):
        # Include a summary.json file to check that we skip reading it.
        with open(os.path.join(dest_dir, "summary.json"), "w") as fh:
            fh.write("dummy_data")
        # Include a *.catapult_json file to check that we skip reading these.
        with open(os.path.join(dest_dir, "foo.catapult_json"), "w") as fh:
            fh.write("dummy_data")

    def WriteExampleDataDir(
        self,
        dir_path,
        mean=1000,
        stddev=100,
        drop_one=False,
        single_boot=False,
        unit="nanoseconds",
    ):
        results = [("ClockGetTimeExample", GenerateTestData(mean, stddev))]
        if not drop_one:
            results.append(("SecondExample", GenerateTestData(2000, 300)))

        if single_boot:
            for test_name, values in results:
                dest_dir = os.path.join(dir_path, "by_boot", "boot0")
                dest_file = os.path.join(
                    dest_dir, "%s.fuchsiaperf.json" % test_name
                )
                if not os.path.exists(dest_dir):
                    os.makedirs(dest_dir)
                    self.AddIgnoredFiles(dest_dir)
                WriteJsonFile(
                    dest_file,
                    [
                        {
                            "label": test_name,
                            "test_suite": "fuchsia.example",
                            "unit": unit,
                            "values": SLOW_INITIAL_RUN + values,
                        }
                    ],
                )
        else:
            for test_name, values in results:
                for idx, value in enumerate(values):
                    dest_dir = os.path.join(
                        dir_path, "by_boot", "boot%06d" % idx
                    )
                    dest_file = os.path.join(
                        dest_dir, "%s.fuchsiaperf.json" % test_name
                    )
                    if not os.path.exists(dest_dir):
                        os.makedirs(dest_dir)
                        self.AddIgnoredFiles(dest_dir)
                    WriteJsonFile(
                        dest_file,
                        [
                            {
                                "label": test_name,
                                "test_suite": "fuchsia.example",
                                "unit": unit,
                                "values": SLOW_INITIAL_RUN + [value],
                            }
                        ],
                    )

    def ExampleDataDir(self, **kwargs):
        dir_path = self.MakeTempDir()
        self.WriteExampleDataDir(dir_path, **kwargs)
        return dir_path

    def test_reading_results_from_dir(self):
        dir_path = self.ExampleDataDir()
        results = perfcompare.StatsFromMultiBootDataset(
            perfcompare.MultiBootDataset(dir_path)
        )
        test_name = "fuchsia.example: ClockGetTimeExample"
        self.assertEqual(
            results[test_name].FormatConfidenceInterval(), "992 +/- 26 ns"
        )

    def test_mean_and_stddev(self):
        values = [10, 5, 15]
        mean_val, stddev_val = perfcompare.MeanAndStddev(values)
        self.assertEqual(mean_val, 10.0)
        self.assertEqual(perfcompare.Mean(values), 10.0)
        self.assertEqual(stddev_val, 5.0)
        # Single-value sample.
        self.assertEqual(perfcompare.MeanAndStddev([123]), (123.0, None))
        # Check error cases.
        self.assertRaises(AssertionError, lambda: perfcompare.Mean([]))
        self.assertRaises(AssertionError, lambda: perfcompare.MeanAndStddev([]))

    # Check that data written using the golden file helper reads back
    # the same.
    def test_golden_file_write_and_read(self):
        temp_file = pathlib.Path(os.path.join(self.MakeTempDir(), "file"))
        writer = GoldenDataOutput()
        writer.AssertCaseEq("a_key", "a_value")
        writer.AssertCaseEq("b_key", "line 1\n" "line 2\n")
        writer.WriteFile(temp_file)
        reader = GoldenDataInput(temp_file)
        reader.AssertCaseEq("a_key", "a_value")
        reader.AssertCaseEq("b_key", "line 1\n" "line 2\n")
        self.assertRaises(
            AssertionError, lambda: reader.AssertCaseEq("a_key", "other_value")
        )

    def test_comparison_no_change(self):
        before_dir = self.ExampleDataDir()
        after_dir = self.ExampleDataDir()
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_no_change", output)

    # Test a regression that is large enough to be flagged.
    def test_comparison_regression(self):
        before_dir = self.ExampleDataDir(mean=1500, stddev=100)
        after_dir = self.ExampleDataDir(mean=1600, stddev=100)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_regression", output)

    # Test an improvement that is large enough to be flagged.
    def test_comparison_improvement(self):
        before_dir = self.ExampleDataDir(mean=1500, stddev=100)
        after_dir = self.ExampleDataDir(mean=1400, stddev=100)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_improvement", output)

    # Test an improvement that is not large enough to be flagged.
    def test_comparison_improvement_small(self):
        before_dir = self.ExampleDataDir(mean=1500, stddev=100)
        after_dir = self.ExampleDataDir(mean=1450, stddev=100)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_improvement_small", output)

    # Test a decrease that is a regression (in contrast to the default
    # where decreases are improvements).
    def test_comparison_regression_biggerisbetter(self):
        unit = "bytes/second"
        before_dir = self.ExampleDataDir(mean=1500, stddev=100, unit=unit)
        after_dir = self.ExampleDataDir(mean=1400, stddev=100, unit=unit)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_regression_biggerisbetter", output)

    # Test an increase that is an improvement (in contrast to the default
    # where increases are regressions).
    def test_comparison_improvement_biggerisbetter(self):
        unit = "bytes/second"
        before_dir = self.ExampleDataDir(mean=1400, stddev=100, unit=unit)
        after_dir = self.ExampleDataDir(mean=1500, stddev=100, unit=unit)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("comparison_improvement_biggerisbetter", output)

    def test_adding_test(self):
        before_dir = self.ExampleDataDir(drop_one=True)
        after_dir = self.ExampleDataDir()
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("adding_test", output)

    def test_removing_test(self):
        before_dir = self.ExampleDataDir()
        after_dir = self.ExampleDataDir(drop_one=True)
        output = ComparePerf([before_dir, after_dir])
        GOLDEN.AssertCaseEq("removing_test", output)

    def test_display_single_dataset(self):
        dataset_dir = self.ExampleDataDir()
        output = ComparePerf([dataset_dir])
        GOLDEN.AssertCaseEq("display_single_dataset", output)

    def test_display_three_datasets(self):
        dataset_dirs = [
            self.ExampleDataDir(mean=1000),
            self.ExampleDataDir(mean=2000, drop_one=True),
            self.ExampleDataDir(mean=3000),
        ]
        output = ComparePerf(dataset_dirs)
        GOLDEN.AssertCaseEq("display_three_datasets", output)

    # Test printing a table of point estimates.
    def test_display_single_boot_single_dataset(self):
        dataset_dir = self.ExampleDataDir(single_boot=True)
        output = ComparePerf([dataset_dir], expected_sample_size=1)
        GOLDEN.AssertCaseEq("display_single_boot_single_dataset", output)

    # Test printing a table of point estimates.
    def test_display_single_boot_two_datasets(self):
        dataset_dirs = [
            self.ExampleDataDir(mean=1000, single_boot=True),
            self.ExampleDataDir(mean=2000, single_boot=True, drop_one=True),
        ]
        output = ComparePerf(dataset_dirs, expected_sample_size=1)
        GOLDEN.AssertCaseEq("display_single_boot_two_datasets", output)

    def test_factor_range_formatting(self):
        # Construct an interval pair of the same type used in the
        # software-under-test, in order to ensure that this test case is
        # realistic.  (Earlier versions used numpy.float64.)  Also check
        # that the interval is well-formed.
        def Interval(min_val, max_val):
            assert min_val <= max_val
            return (float(min_val), float(max_val))

        # Check that the values are of the same type as in the
        # software-under-test.
        interval_test = Interval(10, 20)
        interval_real = perfcompare.Stats([1, 2, 3], "some_unit").interval
        self.assertEqual(type(interval_test[0]), type(interval_real[0]))
        self.assertEqual(type(interval_test[1]), type(interval_real[1]))

        def Format(interval_before, interval_after):
            return perfcompare.FormatFactorRange(
                Interval(*interval_before), Interval(*interval_after)
            )

        self.assertEqual(Format((1, 2), (3, 4)), "1.500-4.000")
        # Test zero "min" values.
        self.assertEqual(Format((0, 2), (3, 4)), "1.500-inf")
        self.assertEqual(Format((1, 2), (0, 4)), "0.000-4.000")
        # Test zero "min" and "max" values.
        self.assertEqual(Format((0, 0), (3, 4)), "inf-inf")
        self.assertEqual(Format((1, 2), (0, 0)), "0.000-0.000")
        # Test zero "max" values, with negative "min".
        self.assertEqual(Format((-1, 0), (3, 4)), "ci_too_wide")
        self.assertEqual(Format((1, 2), (-3, 0)), "ci_too_wide")
        # All values zero.
        self.assertEqual(Format((0, 0), (0, 0)), "no_change")

    def test_mismatch_rate(self):
        self.assertEqual(perfcompare.MismatchRate([(0, 1), (2, 3)]), 1)
        self.assertEqual(perfcompare.MismatchRate([(0, 2), (1, 3)]), 0)
        self.assertEqual(
            perfcompare.MismatchRate([(0, 2), (1, 3), (4, 5)]), 2.0 / 3
        )

    def test_validate_perfcompare(self):
        def MakeExampleDirs(**kwargs):
            by_boot_dir = os.path.join(self.ExampleDataDir(**kwargs), "by_boot")
            return [
                os.path.join(by_boot_dir, name)
                for name in sorted(os.listdir(by_boot_dir))
            ]

        # This is an example input dataset that gives a high mismatch rate,
        # because the data is drawn from two very different distributions.
        results_dirs = MakeExampleDirs(mean=100, stddev=10) + MakeExampleDirs(
            mean=200, stddev=10
        )
        stdout = io.StringIO()
        perfcompare.Main(
            ["validate_perfcompare", "--group_size=5"] + results_dirs, stdout
        )
        output = stdout.getvalue()
        GOLDEN.AssertCaseEq("validate_perfcompare", output)

    def test_biggerisbetter_handling(self):
        self.assertFalse(perfcompare.UnitBiggerIsBetter("ns"))
        self.assertTrue(perfcompare.UnitBiggerIsBetter("bytes/second"))
        self.assertTrue(perfcompare.UnitBiggerIsBetter("frames/second"))
        self.assertTrue(perfcompare.UnitBiggerIsBetter("count_biggerIsBetter"))
        self.assertFalse(
            perfcompare.UnitBiggerIsBetter("count_smallerIsBetter")
        )


class RunLocalTest(TempDirTestCase):
    # Test basic operation of the "run_local" subcommand.
    def test_run_local(self):
        # Destination directory for the full multiboot dataset.  Use a
        # destination path that does not exist yet.
        dest_dir = os.path.join(self.MakeTempDir(), "new_dir")

        # Destination pathnames for process dataset files.
        iter_temp_dir = self.MakeTempDir()
        iter_temp_file = os.path.join(iter_temp_dir, "result.fuchsiaperf.json")
        iter_temp_glob = os.path.join(iter_temp_dir, "*.fuchsiaperf.json")

        data = GenerateData(
            mean=1000,
            stddev_across_boots=10,
            stddev_across_processes=10,
            stddev_across_iters=10,
        )
        commands = []

        # Dummy version of subprocess.check_call() for testing.
        def DummyRunCmd(cmd, shell=False):
            self.assertEqual(shell, True)
            commands.append(cmd)
            if cmd == "set -o errexit -o nounset; my_iter_cmd":
                WriteJsonFile(
                    iter_temp_file,
                    [
                        {
                            "label": "MyTest",
                            "test_suite": "example_suite",
                            "unit": "nanoseconds",
                            "values": data.pop(0)[0],
                        }
                    ],
                )

        stdout = io.StringIO()
        perfcompare.Main(
            [
                "run_local",
                "--boots=4",
                "--iter_file",
                iter_temp_glob,
                "--iter_cmd",
                "my_iter_cmd",
                "--reboot_cmd",
                "my_reboot_cmd",
                "--dest",
                dest_dir,
            ],
            stdout,
            run_cmd=DummyRunCmd,
        )
        self.assertEqual(
            commands,
            [
                "set -o errexit -o nounset; my_reboot_cmd",
                "set -o errexit -o nounset; my_iter_cmd",
            ]
            * 4,
        )
        GOLDEN.AssertCaseEq("run_local", stdout.getvalue())

    # "run_local" should give an error if the temporary files specified by
    # --iter_file already exist.
    def test_error_if_dest_files_already_exist(self):
        dest_dir = os.path.join(self.MakeTempDir(), "new_dir")
        iter_temp_file = os.path.join(
            self.MakeTempDir(), "result.fuchsiaperf.json"
        )
        WriteJsonFile(iter_temp_file, [])
        args = [
            "run_local",
            "--boots=4",
            "--iter_file",
            iter_temp_file,
            "--iter_cmd",
            "my_iter_cmd",
            "--reboot_cmd",
            "my_reboot_cmd",
            "--dest",
            dest_dir,
        ]
        self.assertRaises(
            AssertionError, lambda: perfcompare.Main(args, sys.stdout)
        )

    # Check that error-checking is enabled in the shell commands that
    # run_local runs.
    def test_errexit_error_checking_in_shell_commands(self):
        iter_temp_file = os.path.join(
            self.MakeTempDir(), "result.fuchsiaperf.json"
        )
        stdout = io.StringIO()

        def get_args():
            dest_dir = os.path.join(self.MakeTempDir(), "new_dir")
            return [
                "run_local",
                "--boots=4",
                "--iter_file",
                iter_temp_file,
                "--dest",
                dest_dir,
            ]

        perfcompare.Main(
            get_args() + ["--iter_cmd", "true", "--reboot_cmd", "true"], stdout
        )
        # Check that the failure of the "false" command gets caught.
        self.assertRaises(
            subprocess.CalledProcessError,
            lambda: perfcompare.Main(
                get_args()
                + ["--iter_cmd", "false; true", "--reboot_cmd", "true"],
                stdout,
            ),
        )
        self.assertRaises(
            subprocess.CalledProcessError,
            lambda: perfcompare.Main(
                get_args()
                + ["--iter_cmd", "true", "--reboot_cmd", "false; true"],
                stdout,
            ),
        )


class LocalWorkflowTest(TempDirTestCase):
    def setUp(self):
        super().setUp()
        self.state_file = os.path.join(self.MakeTempDir(), "state.json")
        state_data = {
            "2026-04-09": {
                "launched": {
                    "builder1": {"build_id": "1001"},
                }
            },
            "2026-04-10": {
                "launched": {
                    "builder1": {"build_id": "1002"},
                    "builder2": {"build_id": "1003"},
                }
            },
        }
        with open(self.state_file, "w") as f:
            json.dump(state_data, f)

        # Patch the code for finding the "bb" and "cas" tools so that the
        # test does not depend on the FUCHSIA_DIR env var being set and the
        # tools being present.
        self.enterContext(
            unittest.mock.patch(
                "local_workflow._find_bb_tool", lambda: "bb_for_testing"
            )
        )
        self.enterContext(
            unittest.mock.patch(
                "local_workflow._find_cas_tool", lambda: "cas_for_testing"
            )
        )

    @staticmethod
    def _make_cas_download_fake(
        content: str,
    ) -> Callable[[List[str]], subprocess.CompletedProcess]:
        def handler(cmd: List[str]) -> subprocess.CompletedProcess:
            dir_index = cmd.index("-dir")
            target_dir = cmd[dir_index + 1]
            os.makedirs(target_dir, exist_ok=True)
            with open(os.path.join(target_dir, "metrics.json"), "w") as f:
                f.write(content)
            return subprocess.CompletedProcess(cmd, 0, stdout="download ok")

        return handler

    @staticmethod
    def _make_bb_cas_fake(
        mapping: Dict[Tuple[str, ...], Behavior]
    ) -> Callable[[List[str]], subprocess.CompletedProcess]:
        def side_effect(
            cmd: List[str], **kwargs: Any
        ) -> subprocess.CompletedProcess:
            for patterns, behavior in mapping.items():
                if all(p in cmd for p in patterns):
                    if callable(behavior):
                        return behavior(cmd)
                    elif isinstance(behavior, Exception):
                        raise behavior
                    elif isinstance(behavior, str):
                        return subprocess.CompletedProcess(
                            cmd, 0, stdout=behavior, stderr=""
                        )
                    else:
                        raise ValueError(
                            f"Unexpected behavior type: {type(behavior)}"
                        )
            raise RuntimeError(f"No match found for command: {cmd}")

        return side_effect

    @unittest.mock.patch("subprocess.run")
    def test_download_metrics(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {
                                "cas_instance": "chromium_swarm",
                                "perf_dataset_digest": "123abc/82",
                            }
                        }
                    }
                ),
                ("download",): self._make_cas_download_fake('{"metric": 1.0}'),
            }
        )

        out_dir = self.MakeTempDir()
        stdout = io.StringIO()
        perfcompare.Main(["download_metrics", "1000", out_dir], stdout)

        expected_file = os.path.join(out_dir, "metrics.json")
        self.assertTrue(os.path.exists(expected_file))
        with open(expected_file) as f:
            self.assertEqual(f.read(), '{"metric": 1.0}')

        output = stdout.getvalue()
        self.assertIn("Fetching build info for 1000 using bb", output)
        self.assertIn("Downloading from cas to", output)

    @unittest.mock.patch("subprocess.run")
    def test_download_metrics_perfcompare_digest(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {
                                "cas_instance": "chromium_swarm",
                                "perfcompare_dataset_digest": "123abc/82",
                            }
                        }
                    }
                ),
                ("download",): self._make_cas_download_fake('{"metric": 1.0}'),
            }
        )

        out_dir = self.MakeTempDir()
        stdout = io.StringIO()
        perfcompare.Main(["download_metrics", "1000", out_dir], stdout)

        expected_file = os.path.join(out_dir, "metrics.json")
        self.assertTrue(os.path.exists(expected_file))
        with open(expected_file) as f:
            self.assertEqual(f.read(), '{"metric": 1.0}')

        output = stdout.getvalue()
        self.assertIn("Fetching build info for 1000 using bb", output)
        self.assertIn("Downloading from cas to", output)

    @unittest.mock.patch("subprocess.run")
    def test_download_metrics_missing_properties(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {"cas_instance": "chromium_swarm"}
                        }
                    }
                ),
            }
        )

        with self.assertRaises(KeyError):
            perfcompare.Main(
                ["download_metrics", "1000", self.MakeTempDir()], io.StringIO()
            )

    @unittest.mock.patch("subprocess.run")
    def test_download_metrics_bb_get_fails(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "-p"): subprocess.CalledProcessError(
                    1, ["bb", "get"], stderr="bb error"
                ),
            }
        )

        with self.assertRaises(RuntimeError) as cm:
            perfcompare.Main(
                ["download_metrics", "1000", self.MakeTempDir()], io.StringIO()
            )
        self.assertIn("Download failed for build 1000", str(cm.exception))

    @unittest.mock.patch("subprocess.run")
    def test_download_metrics_cas_download_fails(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {
                                "cas_instance": "chromium_swarm",
                                "perf_dataset_digest": "123abc/82",
                            }
                        }
                    }
                ),
                ("download",): subprocess.CalledProcessError(
                    1, ["cas", "download"], stderr="cas error"
                ),
            }
        )

        with self.assertRaises(RuntimeError) as cm:
            perfcompare.Main(
                ["download_metrics", "1000", self.MakeTempDir()],
                io.StringIO(),
            )
        self.assertIn("Download failed for build 1000", str(cm.exception))

    @unittest.mock.patch("subprocess.run")
    def test_download_baseline_success(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                # Order matters: check for specific 'get' with '-p' before general 'get'
                ("get", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {
                                "cas_instance": "chromium_swarm",
                                "perf_dataset_digest": "123abc/82",
                            }
                        }
                    }
                ),
                ("get",): json.dumps({"status": "SUCCESS"}),
                ("download",): self._make_cas_download_fake('{"metric": 1.0}'),
            }
        )

        out_dir = self.MakeTempDir()
        stdout = io.StringIO()

        perfcompare.Main(
            [
                "download_baseline",
                "--state_file",
                self.state_file,
                "builder2",
                out_dir,
            ],
            stdout,
        )

        expected_file = os.path.join(out_dir, "metrics.json")
        self.assertTrue(os.path.exists(expected_file))
        with open(expected_file) as f:
            self.assertEqual(f.read(), '{"metric": 1.0}')

        output = stdout.getvalue()
        self.assertIn("1003 (SUCCESS) - Triggering download", output)

    @unittest.mock.patch("subprocess.run")
    def test_download_baseline_fallback(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get", "1002"): json.dumps({"status": "FAILURE"}),
                # Order matters: check for specific 'get' with '-p' before general 'get'
                ("get", "1001", "-p"): json.dumps(
                    {
                        "output": {
                            "properties": {
                                "cas_instance": "chromium_swarm",
                                "perf_dataset_digest": "123abc/82",
                            }
                        }
                    }
                ),
                ("get", "1001"): json.dumps({"status": "SUCCESS"}),
                ("download",): self._make_cas_download_fake('{"metric": 1.0}'),
            }
        )

        out_dir = self.MakeTempDir()
        stdout = io.StringIO()

        perfcompare.Main(
            [
                "download_baseline",
                "--state_file",
                self.state_file,
                "builder1",
                out_dir,
            ],
            stdout,
        )

        expected_file = os.path.join(out_dir, "metrics.json")
        self.assertTrue(os.path.exists(expected_file))
        with open(expected_file) as f:
            self.assertEqual(f.read(), '{"metric": 1.0}')

        output = stdout.getvalue()
        self.assertIn("1002 (FAILURE) - Skipping", output)
        self.assertIn("1001 (SUCCESS) - Triggering download", output)

    @unittest.mock.patch("subprocess.run")
    def test_download_baseline_no_successful_build(self, mock_run):
        mock_run.side_effect = self._make_bb_cas_fake(
            {
                ("auth-info",): "auth ok",
                ("get",): json.dumps({"status": "FAILURE"}),
            }
        )

        out_dir = self.MakeTempDir()
        stdout = io.StringIO()

        perfcompare.Main(
            [
                "download_baseline",
                "--state_file",
                self.state_file,
                "builder1",
                out_dir,
            ],
            stdout,
        )

        output = stdout.getvalue()
        self.assertIn("1002 (FAILURE) - Skipping", output)
        self.assertIn("1001 (FAILURE) - Skipping", output)
        self.assertIn("No successful sample build IDs found", output)

    def test_extract_metrics(self):
        data_dir = self.MakeTempDir()
        WriteJsonFile(
            os.path.join(data_dir, "test.fuchsiaperf.json"),
            [
                {
                    "test_suite": "fuchsia.test",
                    "label": "metric1",
                    "values": [10.5],
                },
                {
                    "test_suite": "fuchsia.test",
                    "label": "metric2",
                    "values": [20.0],
                },
                {
                    "test_suite": "other.suite",
                    "label": "metric3",
                    "values": [30.0],
                },
            ],
        )

        stdout = io.StringIO()
        perfcompare.Main(["extract_metrics", "fuchsia", data_dir], stdout)

        output = stdout.getvalue()
        self.assertIn("fuchsia.test,metric1,10.5", output)
        self.assertIn("fuchsia.test,metric2,20.0", output)
        self.assertNotIn("other.suite", output)

    def test_extract_metrics_aggregation(self):
        data_dir = self.MakeTempDir()
        WriteJsonFile(
            os.path.join(data_dir, "test1.fuchsiaperf.json"),
            [
                {
                    "test_suite": "fuchsia.test",
                    "label": "metric1",
                    "values": [10.5],
                }
            ],
        )
        WriteJsonFile(
            os.path.join(data_dir, "test2.fuchsiaperf.json"),
            [
                {
                    "test_suite": "fuchsia.test",
                    "label": "metric1",
                    "values": [15.0],
                }
            ],
        )

        stdout = io.StringIO()
        perfcompare.Main(["extract_metrics", "fuchsia", data_dir], stdout)

        stdout.seek(0)
        reader = csv.reader(stdout)
        rows = list(reader)
        self.assertEqual(rows, [["fuchsia.test", "metric1", "10.5", "15.0"]])

    def test_extract_metrics_empty_values(self):
        data_dir = self.MakeTempDir()
        WriteJsonFile(
            os.path.join(data_dir, "test.fuchsiaperf.json"),
            [
                {
                    "test_suite": "fuchsia.test",
                    "label": "metric1",
                    "values": [],
                }
            ],
        )

        with self.assertRaises(statistics.StatisticsError):
            perfcompare.Main(
                ["extract_metrics", "fuchsia", data_dir], io.StringIO()
            )

    def test_extract_metrics_malformed_json(self):
        data_dir = self.MakeTempDir()
        file_path = os.path.join(data_dir, "bad.fuchsiaperf.json")
        with open(file_path, "w") as f:
            f.write("{invalid}")

        with self.assertRaises(RuntimeError) as cm:
            perfcompare.Main(
                ["extract_metrics", "fuchsia", data_dir], io.StringIO()
            )
        self.assertIn("Malformed JSON", str(cm.exception))


if __name__ == "__main__":
    TestMain()