def test_load_dataset_is_mnist_called(client_task_definition_data):
client = Client()
client._load_mnist_dataset = MagicMock()
client._cluster_request_data = client_task_definition_data
client.load_dataset()
client._load_mnist_dataset.assert_called_once()
def test_load_csv_dataset_from_hardrive(client_task_definition_csv_data):
client = Client()
client._load_mnist_dataset = MagicMock()
client._cluster_request_data = client_task_definition_csv_data
training_data = client._load_csv_dataset()
assert type(training_data) == mxnet.gluon.data.DataLoader
for feature_batch, label_batch in training_data:
assert len(feature_batch) == 100
def test_load_dataset_is_csv_called(client_task_definition_data):
client = Client()
client._load_mnist_dataset = MagicMock()
client._load_csv_dataset = MagicMock()
client._cluster_request_data = client_task_definition_data
client._cluster_request_data['ml']['dataset']['format'] = 'CSV'
client.load_dataset()
client._load_csv_dataset.assert_called_once()