vedithal-amd
418 righe
15 KiB
Python
418 righe
15 KiB
Python
##############################################################################bl
|
|
# MIT License
|
|
#
|
|
# Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
##############################################################################el
|
|
|
|
|
|
import logging
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, Mock, call, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
logging.TRACE = logging.DEBUG - 5
|
|
logging.addLevelName(logging.TRACE, "TRACE")
|
|
|
|
|
|
def trace_logger(message, *args, **kwargs):
|
|
logging.log(logging.TRACE, message, *args, **kwargs)
|
|
|
|
|
|
setattr(logging, "trace", trace_logger)
|
|
|
|
from db_connector import DatabaseConnector
|
|
|
|
"""
|
|
Tests for the DatabaseConnector class that tests almost methods with initialization,
|
|
CSV import, database removal, and error handling.
|
|
The tests use mocks instead of a real MongoDB server for speed and reliability.
|
|
"""
|
|
|
|
|
|
class TestDatabaseConnector:
|
|
@pytest.fixture
|
|
def mock_args_import(self):
|
|
"""Mock arguments for import operation"""
|
|
args = Mock()
|
|
args.username = "test_user"
|
|
args.password = "test_pass"
|
|
args.host = "localhost"
|
|
args.port = 27017
|
|
args.team = "test_team"
|
|
args.workload = "/app/tests/workloads/device_filter/MI100"
|
|
args.upload = True
|
|
args.remove = False
|
|
args.kernel_verbose = False
|
|
return args
|
|
|
|
@pytest.fixture
|
|
def mock_args_remove(self):
|
|
"""Mock arguments for remove operation"""
|
|
args = Mock()
|
|
args.username = "test_user"
|
|
args.password = "test_pass"
|
|
args.host = "localhost"
|
|
args.port = 27017
|
|
args.team = "test_team"
|
|
args.workload = "rocprofiler-compute_test_team_workload_mi100"
|
|
args.upload = False
|
|
args.remove = True
|
|
args.kernel_verbose = False
|
|
return args
|
|
|
|
def test_init(self, mock_args_import):
|
|
"""Test DatabaseConnector initialization"""
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
assert connector.args == mock_args_import
|
|
assert isinstance(connector.cache, dict)
|
|
assert len(connector.cache) == 0
|
|
|
|
expected_connection_info = {
|
|
"username": "test_user",
|
|
"password": "test_pass",
|
|
"host": "localhost",
|
|
"port": "27017",
|
|
"team": "test_team",
|
|
"workload": "/app/tests/workloads/device_filter/MI100",
|
|
"db": None,
|
|
}
|
|
assert connector.connection_info == expected_connection_info
|
|
assert connector.interaction_type is None
|
|
assert connector.client is None
|
|
|
|
@patch("db_connector.pd.read_csv")
|
|
@patch("db_connector.Path")
|
|
def test_prep_import_success(self, mock_path, mock_read_csv, mock_args_import):
|
|
"""Test successful prep_import"""
|
|
# Setup mocks
|
|
mock_path.return_value.joinpath.return_value = "/fake/path/sysinfo.csv"
|
|
mock_path.return_value.is_file.return_value = True
|
|
|
|
mock_sysinfo = pd.DataFrame(
|
|
{"gpu_model": ["MI100 "], "workload_name": [" test_workload"]}
|
|
)
|
|
mock_read_csv.return_value = mock_sysinfo
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
connector.prep_import()
|
|
|
|
expected_db = "rocprofiler-compute_test_team_test_workload_MI100"
|
|
assert connector.connection_info["db"] == expected_db
|
|
|
|
@patch("db_connector.pd.read_csv")
|
|
@patch("db_connector.Path")
|
|
def test_prep_import_missing_file(self, mock_path, mock_read_csv, mock_args_import):
|
|
"""Test prep_import when sysinfo.csv is missing"""
|
|
mock_path.return_value.joinpath.return_value = "/fake/path/sysinfo.csv"
|
|
mock_path.return_value.is_file.return_value = False
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch(
|
|
"db_connector.console_error", side_effect=SystemExit(1)
|
|
) as mock_console_error:
|
|
with pytest.raises(SystemExit):
|
|
connector.prep_import()
|
|
|
|
mock_console_error.assert_called_with(
|
|
"database", "Unable to parse SoC and/or workload name from sysinfo.csv"
|
|
)
|
|
|
|
@patch("db_connector.pd.read_csv")
|
|
@patch("db_connector.Path")
|
|
def test_prep_import_key_error(self, mock_path, mock_read_csv, mock_args_import):
|
|
"""Test prep_import when required fields are missing"""
|
|
mock_path.return_value.joinpath.return_value = "/fake/path/sysinfo.csv"
|
|
mock_path.return_value.is_file.return_value = True
|
|
|
|
mock_sysinfo = pd.DataFrame({"other_column": ["value"]})
|
|
mock_read_csv.return_value = mock_sysinfo
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch(
|
|
"db_connector.console_error", side_effect=SystemExit(1)
|
|
) as mock_console_error:
|
|
with pytest.raises(SystemExit):
|
|
connector.prep_import()
|
|
|
|
assert mock_console_error.called
|
|
error_call = mock_console_error.call_args[0][0]
|
|
assert "Outdated workload" in error_call
|
|
|
|
@patch("db_connector.tqdm")
|
|
@patch("db_connector.os.listdir")
|
|
@patch("db_connector.console_log")
|
|
@patch("db_connector.console_warning")
|
|
@patch("db_connector.kernel_name_shortener")
|
|
@patch("db_connector.MongoClient")
|
|
@patch("db_connector.pd.read_csv")
|
|
def test_db_import_success(
|
|
self,
|
|
mock_read_csv,
|
|
mock_mongo_client,
|
|
mock_kernel_shortener,
|
|
mock_console_warning,
|
|
mock_console_log,
|
|
mock_listdir,
|
|
mock_tqdm,
|
|
mock_args_import,
|
|
):
|
|
"""Test successful database import"""
|
|
mock_listdir.return_value = ["test_data.csv", "empty_file.csv", "non_csv.txt"]
|
|
mock_tqdm.return_value = mock_listdir.return_value
|
|
|
|
test_df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
|
|
mock_read_csv.side_effect = [test_df, pd.errors.EmptyDataError()]
|
|
|
|
mock_client_instance = MagicMock()
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_workload_db = MagicMock()
|
|
mock_workload_col = MagicMock()
|
|
|
|
mock_mongo_client.return_value = mock_client_instance
|
|
mock_client_instance.__getitem__.side_effect = lambda x: {
|
|
"rocprofiler-compute_test_team_test_workload_MI100": mock_db,
|
|
"workload_names": mock_workload_db,
|
|
}.get(x, mock_db)
|
|
mock_db.__getitem__.return_value = mock_collection
|
|
mock_workload_db.__getitem__.return_value = mock_workload_col
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
connector.connection_info["workload"] = "/fake/workload/path"
|
|
connector.client = mock_client_instance
|
|
|
|
with patch.object(connector, "prep_import") as mock_prep:
|
|
mock_prep.return_value = None
|
|
connector.connection_info[
|
|
"db"
|
|
] = "rocprofiler-compute_test_team_test_workload_MI100"
|
|
|
|
connector.db_import()
|
|
|
|
mock_collection.insert_many.assert_called_once()
|
|
mock_workload_col.replace_one.assert_called_once()
|
|
|
|
@patch("db_connector.console_log")
|
|
def test_db_remove_success(self, mock_console_log, mock_args_remove):
|
|
"""Test successful database removal"""
|
|
mock_client = MagicMock()
|
|
mock_db_to_remove = MagicMock()
|
|
mock_workload_names_db = MagicMock()
|
|
mock_names_col = MagicMock()
|
|
|
|
mock_client.__getitem__.side_effect = lambda x: {
|
|
"rocprofiler-compute_test_team_workload_mi100": mock_db_to_remove,
|
|
"workload_names": mock_workload_names_db,
|
|
}[x]
|
|
mock_workload_names_db.__getitem__.return_value = mock_names_col
|
|
mock_db_to_remove.list_collection_names.return_value = ["col1", "col2"]
|
|
|
|
connector = DatabaseConnector(mock_args_remove)
|
|
connector.client = mock_client
|
|
|
|
connector.db_remove()
|
|
|
|
mock_client.drop_database.assert_called_once_with(mock_db_to_remove)
|
|
mock_names_col.delete_many.assert_called_once_with(
|
|
{"name": "rocprofiler-compute_test_team_workload_mi100"}
|
|
)
|
|
|
|
def test_pre_processing_no_action_specified(self, mock_args_import):
|
|
"""Test pre_processing when neither upload nor remove is specified"""
|
|
mock_args_import.upload = False
|
|
mock_args_import.remove = False
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
def test_pre_processing_remove_invalid_workload_name(self, mock_args_remove):
|
|
"""Test pre_processing remove with invalid workload name"""
|
|
mock_args_remove.workload = "invalid_name"
|
|
|
|
connector = DatabaseConnector(mock_args_remove)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
def test_pre_processing_remove_missing_host_username(self, mock_args_remove):
|
|
"""Test pre_processing remove with missing host/username"""
|
|
mock_args_remove.host = None
|
|
mock_args_remove.username = None
|
|
|
|
connector = DatabaseConnector(mock_args_remove)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
def test_pre_processing_remove_protected_database(self, mock_args_remove):
|
|
"""Test pre_processing remove with protected database names"""
|
|
mock_args_remove.workload = "admin"
|
|
|
|
connector = DatabaseConnector(mock_args_remove)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
@patch("db_connector.Path")
|
|
@patch("db_connector.is_workload_empty")
|
|
@patch("db_connector.getpass.getpass")
|
|
@patch("db_connector.console_log")
|
|
@patch("db_connector.MongoClient")
|
|
def test_pre_processing_import_password_prompt_success(
|
|
self,
|
|
mock_mongo_client,
|
|
mock_console_log,
|
|
mock_getpass,
|
|
mock_is_workload_empty,
|
|
mock_path,
|
|
mock_args_import,
|
|
):
|
|
"""Test pre_processing import with password prompt success"""
|
|
mock_args_import.password = ""
|
|
mock_getpass.return_value = "prompted_password"
|
|
|
|
mock_path.return_value.absolute.return_value.is_dir.return_value = True
|
|
mock_path.return_value.absolute.return_value.resolve.return_value = (
|
|
"/resolved/path"
|
|
)
|
|
|
|
mock_client_instance = MagicMock()
|
|
mock_mongo_client.return_value = mock_client_instance
|
|
mock_client_instance.server_info.return_value = {}
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
connector.pre_processing()
|
|
|
|
mock_getpass.assert_called_once()
|
|
mock_console_log.assert_called_with("database", "Password received")
|
|
|
|
@patch("db_connector.Path")
|
|
@patch("db_connector.is_workload_empty")
|
|
@patch("db_connector.MongoClient")
|
|
def test_pre_processing_import_connection_failure(
|
|
self, mock_mongo_client, mock_is_workload_empty, mock_path, mock_args_import
|
|
):
|
|
"""Test pre_processing import with MongoDB connection failure"""
|
|
mock_path.return_value.absolute.return_value.is_dir.return_value = True
|
|
mock_path.return_value.absolute.return_value.resolve.return_value = (
|
|
"/resolved/path"
|
|
)
|
|
|
|
mock_client_instance = MagicMock()
|
|
mock_mongo_client.return_value = mock_client_instance
|
|
mock_client_instance.server_info.side_effect = Exception("Connection failed")
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
@patch("db_connector.Path")
|
|
@patch("db_connector.is_workload_empty")
|
|
def test_pre_processing_import_missing_required_fields(
|
|
self, mock_is_workload_empty, mock_path, mock_args_import
|
|
):
|
|
"""Test pre_processing import with missing required fields"""
|
|
mock_args_import.host = None
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
@patch("db_connector.Path")
|
|
def test_pre_processing_import_invalid_workload_path(
|
|
self, mock_path, mock_args_import
|
|
):
|
|
"""Test pre_processing import with invalid workload path"""
|
|
mock_path.return_value.absolute.return_value.is_dir.return_value = False
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
def test_pre_processing_import_team_name_too_long(self, mock_args_import):
|
|
"""Test pre_processing import with team name exceeding limit"""
|
|
mock_args_import.team = "this_team_name_is_way_too_long"
|
|
|
|
connector = DatabaseConnector(mock_args_import)
|
|
|
|
with patch("db_connector.console_error", side_effect=SystemExit(1)):
|
|
with pytest.raises(SystemExit):
|
|
connector.pre_processing()
|
|
|
|
|
|
class TestDatabaseConnectorIntegration:
|
|
"""Simple integration test"""
|
|
|
|
@patch("db_connector.Path")
|
|
@patch("db_connector.pd.read_csv")
|
|
def test_prep_import_with_real_workload_path(self, mock_read_csv, mock_path):
|
|
"""Test prep_import with actual workload path structure"""
|
|
args = Mock()
|
|
args.username = "test_user"
|
|
args.password = "test_pass"
|
|
args.host = "localhost"
|
|
args.port = 27017
|
|
args.team = "test_team"
|
|
args.workload = "/app/tests/workloads/device_filter/MI100"
|
|
args.upload = True
|
|
args.remove = False
|
|
args.kernel_verbose = False
|
|
|
|
mock_path.return_value.joinpath.return_value = (
|
|
"/app/tests/workloads/device_filter/MI100/sysinfo.csv"
|
|
)
|
|
mock_path.return_value.is_file.return_value = True
|
|
|
|
mock_sysinfo = pd.DataFrame(
|
|
{"gpu_model": ["MI100"], "workload_name": ["device_filter"]}
|
|
)
|
|
mock_read_csv.return_value = mock_sysinfo
|
|
|
|
connector = DatabaseConnector(args)
|
|
connector.prep_import()
|
|
|
|
expected_db = "rocprofiler-compute_test_team_device_filter_MI100"
|
|
assert connector.connection_info["db"] == expected_db
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|