tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

add mocks for loading datasets in cli train tests (#2497) [skip ci]

Browse Source 
* add mocks for loading datasets in cli train tests

* Apply suggestions from code review to fix patched module for preprocess

Co-authored-by: NanoCode012 <nano@axolotl.ai>

---------

Co-authored-by: NanoCode012 <nano@axolotl.ai>
This commit is contained in:
Wing Lian
2025-04-11 09:51:59 -04:00
committed by GitHub
parent 22c562533d
commit 7e7180fa10
2 changed files with 60 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
tests/cli/test_cli_preprocess.py
View File

@@ -2,7 +2,7 @@
import shutil import shutil
from pathlib import Path from pathlib import Path
from unittest.mock import patch from unittest.mock import MagicMock, patch
import pytest import pytest
@@ -26,12 +26,15 @@ def test_preprocess_config_not_found(cli_runner):
def test_preprocess_basic(cli_runner, config_path): def test_preprocess_basic(cli_runner, config_path):
"""Test basic preprocessing with minimal config""" """Test basic preprocessing with minimal config"""
with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli: with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli:
result = cli_runner.invoke(cli, ["preprocess", str(config_path)]) with patch("axolotl.cli.preprocess.load_datasets") as mock_load_datasets:
assert result.exit_code == 0 mock_load_datasets.return_value = MagicMock()
mock_do_cli.assert_called_once() result = cli_runner.invoke(cli, ["preprocess", str(config_path)])
assert mock_do_cli.call_args.kwargs["config"] == str(config_path) assert result.exit_code == 0
assert mock_do_cli.call_args.kwargs["download"] is True
mock_do_cli.assert_called_once()
assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
assert mock_do_cli.call_args.kwargs["download"] is True
def test_preprocess_without_download(cli_runner, config_path): def test_preprocess_without_download(cli_runner, config_path):
@@ -54,19 +57,22 @@ def test_preprocess_custom_path(cli_runner, tmp_path, valid_test_config):
config_path.write_text(valid_test_config) config_path.write_text(valid_test_config)
with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli: with patch("axolotl.cli.preprocess.do_cli") as mock_do_cli:
result = cli_runner.invoke( with patch("axolotl.cli.preprocess.load_datasets") as mock_load_datasets:
cli, mock_load_datasets.return_value = MagicMock()
[
"preprocess",
str(config_path),
"--dataset-prepared-path",
str(custom_path.absolute()),
],
)
assert result.exit_code == 0
mock_do_cli.assert_called_once() result = cli_runner.invoke(
assert mock_do_cli.call_args.kwargs["config"] == str(config_path) cli,
assert mock_do_cli.call_args.kwargs["dataset_prepared_path"] == str( [
custom_path.absolute() "preprocess",
) str(config_path),
"--dataset-prepared-path",
str(custom_path.absolute()),
],
)
assert result.exit_code == 0
mock_do_cli.assert_called_once()
assert mock_do_cli.call_args.kwargs["config"] == str(config_path)
assert mock_do_cli.call_args.kwargs["dataset_prepared_path"] == str(
custom_path.absolute()
)

62
tests/cli/test_cli_train.py
View File

@@ -29,19 +29,21 @@ class TestTrainCommand(BaseCliTest):
with patch("axolotl.cli.train.train") as mock_train: with patch("axolotl.cli.train.train") as mock_train:
mock_train.return_value = (MagicMock(), MagicMock(), MagicMock()) mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
with patch("axolotl.cli.train.load_datasets") as mock_load_datasets:
mock_load_datasets.return_value = MagicMock()
result = cli_runner.invoke( result = cli_runner.invoke(
cli, cli,
[ [
"train", "train",
str(config_path), str(config_path),
"--no-accelerate", "--no-accelerate",
], ],
catch_exceptions=False, catch_exceptions=False,
) )
assert result.exit_code == 0 assert result.exit_code == 0
mock_train.assert_called_once() mock_train.assert_called_once()
def test_train_cli_overrides(self, cli_runner, tmp_path, valid_test_config): def test_train_cli_overrides(self, cli_runner, tmp_path, valid_test_config):
"""Test CLI arguments properly override config values""" """Test CLI arguments properly override config values"""
@@ -49,23 +51,25 @@ class TestTrainCommand(BaseCliTest):
with patch("axolotl.cli.train.train") as mock_train: with patch("axolotl.cli.train.train") as mock_train:
mock_train.return_value = (MagicMock(), MagicMock(), MagicMock()) mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
with patch("axolotl.cli.train.load_datasets") as mock_load_datasets:
mock_load_datasets.return_value = MagicMock()
result = cli_runner.invoke( result = cli_runner.invoke(
cli, cli,
[ [
"train", "train",
str(config_path), str(config_path),
"--learning-rate", "--learning-rate",
"1e-4", "1e-4",
"--micro-batch-size", "--micro-batch-size",
"2", "2",
"--no-accelerate", "--no-accelerate",
], ],
catch_exceptions=False, catch_exceptions=False,
) )
assert result.exit_code == 0 assert result.exit_code == 0
mock_train.assert_called_once() mock_train.assert_called_once()
cfg = mock_train.call_args[1]["cfg"] cfg = mock_train.call_args[1]["cfg"]
assert cfg["learning_rate"] == 1e-4 assert cfg["learning_rate"] == 1e-4
assert cfg["micro_batch_size"] == 2 assert cfg["micro_batch_size"] == 2