Add tests using mocker to main.py

src/test_main.py

@@ -0,0 +1,31 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from main import Net, trainset, transform
+
+
+def test_data_loading_and_preprocessing():
+    """
+    Test the data loading and preprocessing steps.
+    This test asserts that the correct transformations are applied to the dataset and that the DataLoader is created with the correct parameters.
+    """
+    with patch("main.datasets.MNIST", new=Mock()) as mock_dataset, patch(
+        "main.DataLoader", new=Mock()
+    ) as mock_dataloader:
+        pytest.assertTrue(mock_dataset.called)
+        pytest.assertEqual(mock_dataset.call_args[1]["transform"], transform)
+        pytest.assertTrue(mock_dataloader.called)
+        pytest.assertEqual(mock_dataloader.call_args[0][0], trainset)
+        pytest.assertEqual(mock_dataloader.call_args[1]["batch_size"], 64)
+        pytest.assertTrue(mock_dataloader.call_args[1]["shuffle"])
+
+
+def test_model_definition():
+    """
+    Test the model definition.
+    This test asserts that the model is defined correctly.
+    """
+    with patch("main.Net", new=Mock()) as mock_model:
+        pytest.assertTrue(mock_model.called)
+        pytest.assertIsInstance(mock_model.return_value, Net)