Merge pull request #25 from ngmisl/sweep/feature/add-unit-tests

Add Unit Tests for Main Functions

test_main.py

@@ -0,0 +1,57 @@
+import unittest
+from unittest.mock import patch
+import main
+
+class TestMainFunctions(unittest.TestCase):
+    @patch('main.Web3')
+    def test_main(self, mock_web3):
+        # Set up test data
+        mock_web3.eth.contract.return_value = 'test_contract'
+
+        # Call the function with test inputs
+        result = main.main('test_contract')
+
+        # Compare actual output with expected output using assertEqual
+        self.assertEqual(result, 'expected_result')
+
+    @patch('main.Web3')
+    def test_total_supply(self, mock_web3):
+        # Set up test data
+        mock_web3.eth.contract.return_value = 'test_contract'
+
+        # Call the function with test inputs
+        result = main.total_supply('test_contract')
+
+        # Compare actual output with expected output using assertEqual
+        self.assertEqual(result, 'expected_result')
+
+    @patch('main.Web3')
+    def test_write_to_csv(self, mock_web3):
+        # Set up test data
+        token_owners = {'1': 'wallet1', '2': 'wallet2'}
+        file_name = 'test.csv'
+
+        # Call the function with test inputs
+        main.write_to_csv(token_owners, file_name)
+
+        # Check if the file was written correctly
+        with open(file_name, 'r') as f:
+            content = f.read()
+        self.assertEqual(content, 'expected_content')
+
+    @patch('main.Web3')
+    def test_create_clean_snapshot(self, mock_web3):
+        # Set up test data
+        input_file = 'test_input.csv'
+        output_file = 'test_output.csv'
+
+        # Call the function with test inputs
+        main.create_clean_snapshot(input_file, output_file)
+
+        # Check if the file was written correctly
+        with open(output_file, 'r') as f:
+            content = f.read()
+        self.assertEqual(content, 'expected_content')
+
+if __name__ == '__main__':
+    unittest.main()