Add tests using mocker to main.py

src/main.py

@@ -18,6 +18,14 @@
 # Step 2: Define the PyTorch Model
 class Net(nn.Module):
     def __init__(self):
+        """
+        Initializes the Net class which is a subclass of the nn.Module class.
+        Defines the architecture of the neural network with three fully connected layers (fc1, fc2, fc3).
+
+        Inputs: None apart from self.
+        Outputs: None.
+        Side effects: Modifies the state of the Net instance by setting the fc1, fc2, and fc3 attributes.
+        """
         super().__init__()
         self.fc1 = nn.Linear(28 * 28, 128)
         self.fc2 = nn.Linear(128, 64)

tests/test_main.py

@@ -0,0 +1,21 @@
+import pytest
+from unittest.mock import Mock, patch
+from src import main
+
+def test_transform():
+    mock_image = Mock(spec=Image.Image)
+    tensor = main.transform(mock_image)
+    assert isinstance(tensor, torch.Tensor)
+    assert tensor.shape == (1, 28, 28)
+    assert tensor.min() >= -1
+    assert tensor.max() <= 1
+
+@patch('torchvision.datasets.MNIST')
+@patch('torch.utils.data.DataLoader')
+def test_trainloader(mock_dataloader, mock_mnist):
+    mock_mnist.return_value = Mock()
+    mock_dataloader.return_value = Mock()
+    trainloader = main.trainloader
+    mock_mnist.assert_called_once_with('.', download=True, train=True, transform=main.transform)
+    mock_dataloader.assert_called_once_with(mock_mnist.return_value, batch_size=64, shuffle=True)
+    assert trainloader == mock_dataloader.return_value