clean test

Author: aysim319

nssp/tests/test_pull.py

@@ -1,29 +1,17 @@
-from datetime import datetime, date
 import json
-import unittest
 from unittest.mock import patch, MagicMock
-import tempfile
 import os
-import time
-from datetime import datetime
-import pdb
-import pandas as pd
-import pandas.api.types as ptypes
 
 from delphi_nssp.pull import (
     pull_nssp_data,
 )
-from delphi_nssp.constants import (
-    SIGNALS,
-    NEWLINE,
-    SIGNALS_MAP,
-    TYPE_DICT,
-)
+from delphi_nssp.constants import SIGNALS
 
+from delphi_utils import get_structured_logger
 
-class TestPullNSSPData(unittest.TestCase):
+class TestPullNSSPData:
     @patch("delphi_nssp.pull.Socrata")
-    def test_pull_nssp_data(self, mock_socrata):
+    def test_pull_nssp_data(self, mock_socrata, caplog):
         # Load test data
         with open("test_data/page.txt", "r") as f:
             test_data = json.load(f)
@@ -40,19 +28,18 @@ def test_pull_nssp_data(self, mock_socrata):
 
         custom_run = False
 
-        mock_logger = MagicMock()
+        logger = get_structured_logger()
 
         # Call function with test token
         test_token = "test_token"
-        result = pull_nssp_data(test_token, backup_dir, custom_run, mock_logger)
-        print(result)
+        result = pull_nssp_data(test_token, backup_dir, custom_run, logger)
 
         # Check logger used:
-        mock_logger.info.assert_called()
+        assert "Backup file created" in caplog.text
 
         # Check that backup file was created
         backup_files = os.listdir(backup_dir)
-        assert len(backup_files) == 1, "Backup file was not created"
+        assert len(backup_files) == 2, "Backup file was not created"
 
         # Check that Socrata client was initialized with correct arguments
         mock_socrata.assert_called_once_with("data.cdc.gov", test_token)
@@ -71,6 +58,3 @@ def test_pull_nssp_data(self, mock_socrata):
         for signal in SIGNALS:
             assert result[signal].notnull().all(), f"{signal} has rogue NaN"
 
-
-if __name__ == "__main__":
-    unittest.main()