From 7dbf97d23084fc7bce7ce9193af23eeb6eefaf6b Mon Sep 17 00:00:00 2001 From: Fahim Imaduddin Dalvi Date: Fri, 8 Sep 2023 15:53:56 +0300 Subject: [PATCH] Add test for multiconfig assets --- tests/test_benchmark.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 2bb35f46..7d86fc4f 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -67,6 +67,15 @@ def prompt(input_sample): raise Exception("Fail!") +class MockMultiConfigAsset(MockAsset): + @staticmethod + def config(): + return [ + {"name": "Subasset 1", "config": MockAsset.config()}, + {"name": "Subasset 2", "config": MockAsset.config()}, + ] + + @patch("llmebench.utils.import_source_file", MagicMock(return_value=MockAsset)) class TestBenchmarkAssetFinder(unittest.TestCase): @classmethod @@ -270,3 +279,24 @@ def test_multiple_assets_with_failure(self, asset_finder_mock): with open(Path(self.results_dir.name) / "all_results.json") as fp: results = json.load(fp) self.assertEqual(len(results), 2) + + @patch("llmebench.utils.import_source_file") + def test_multi_config_asset(self, asset_importer_mock): + "Run benchmark with multiconfig asset" + + # Create dummy asset file + (Path(self.benchmark_dir.name) / "sample.py").touch(exist_ok=True) + + asset_importer_mock.return_value = MockMultiConfigAsset + + testargs = ["llmebench", self.benchmark_dir.name, self.results_dir.name] + with patch.object(sys, "argv", testargs): + llmebench.benchmark.main() + + with open(Path(self.results_dir.name) / "all_results.json") as fp: + results = json.load(fp) + config = MockMultiConfigAsset.config() + self.assertEqual(len(results), len(config)) + + for subconfig in config: + self.assertIn(f"sample/{subconfig['name']}", results)