Add tests using mocker to main.py

requirements.txt

@@ -33,7 +33,7 @@ starlette==0.27.0
 sympy==1.12
 tomli==2.0.1
 tomlkit==0.12.1
-torch==2.1.0+cpu
+torch==1.8.1
 torchaudio==2.1.0+cpu
 torchvision==0.16.0+cpu
 typing_extensions==4.8.0

src/test_main.py

@@ -0,0 +1,30 @@
+import unittest
+from unittest.mock import patch
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from main import Net, transform
+
+class TestMain(unittest.TestCase):
+    def setUp(self):
+        self.model = Net()
+        self.transform = transform
+
+    def test_Net(self):
+        # Test the forward pass
+        input_tensor = torch.randn(1, 1, 28, 28)
+        output = self.model(input_tensor)
+        self.assertEqual(output.size(), (1, 10))
+
+    @patch('torchvision.datasets.MNIST')
+    @patch('torch.utils.data.DataLoader')
+    def test_data_loading(self, mock_dataloader, mock_dataset):
+        # Mock the MNIST dataset
+        mock_dataset.return_value = datasets.MNIST('.', download=True, train=True, transform=self.transform)
+        # Mock the DataLoader
+        mock_dataloader.return_value = DataLoader(mock_dataset, batch_size=64, shuffle=True)
+        # Assert that the DataLoader is called with the correct arguments
+        mock_dataloader.assert_called_with(mock_dataset, batch_size=64, shuffle=True)
+
+if __name__ == '__main__':
+    unittest.main()