refactored tests for the data extraction module and added Makefile

Author: jibbs1703

Makefile

@@ -0,0 +1,17 @@
+# Default commit message
+COMMIT_MSG ?= "refactored tests for the data extraction module and added Makefile"
+
+# Target to add changes to staging
+add:
+	   git add .
+
+# Target to commit with a message
+commit: add
+	    git commit -m $(COMMIT_MSG)
+
+# Target to push to the default branch
+push: commit
+	     git push
+
+# Run all targets
+all: push

main.py

@@ -2,7 +2,7 @@
 
 import argparse
 
-from src.extract.extract import run_extract
+from src.extract.extract import run_extraction
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Connect to MongoDB.")
@@ -20,12 +20,12 @@
         "--collections",
         type=str,
         nargs="+",
-        default=["users"],
-        help="MongoDB collection names",
+        default=["category", "events", "listing", "sales", "users", "venue"],
+        help="Tickit DB collection names",
     )
     args = parser.parse_args()
 
-    run_extract(
+    run_extraction(
         host=args.host,
         port=args.port,
         database=args.database,

src/extract/extract.py

@@ -46,7 +46,7 @@ def get_data_from_collection(db: Database, collection_name: str) -> tuple[String
     return StringIO(json_data), f"raw-files/{collection_name}.json"
 
 
-def run_extract(host: str, port: int, database: str, collections: list[str], bucket_name: str):
+def run_extraction(host: str, port: int, database: str, collections: list[str], bucket_name: str):
     """Run the Extraction Module Logic"""
     db = connect_to_mongodb(host, port, database)
     for collection in collections:

tests/conftest.py

@@ -1,9 +1,27 @@
 """Fixtures for module tests."""
 
+from unittest.mock import MagicMock, patch
+
 import pytest
+from pymongo.collection import Collection
+from pymongo.database import Database
 
 
 @pytest.fixture
 def mock_connection(mocker):
     """Fixture to mock MongoClient."""
     return mocker.patch("src.extract.extract.MongoClient")
+
+
+@pytest.fixture
+def mock_logger():
+    with patch("src.extract.extract.logger") as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_db():
+    mock = MagicMock(spec=Database)
+    mock_collection = MagicMock(spec=Collection)
+    mock.__getitem__.return_value = mock_collection
+    return mock, mock_collection

tests/test_extract.py

@@ -1,11 +1,13 @@
 """Tests for the extract module."""
 
+import json
+from io import StringIO
 from unittest.mock import MagicMock
 
 import pytest
 from pymongo.errors import ConnectionFailure
 
-from src.extract.extract import connect_to_mongodb
+from src.extract.extract import connect_to_mongodb, get_data_from_collection
 
 
 @pytest.mark.unit
@@ -38,3 +40,49 @@ def test_failed_connection(mock_connection):
 
     # Ensure result is None when connection fails
     assert result is None
+
+
+def test_get_data_from_collection_success(mock_db, mock_logger):
+    mock_db_instance, mock_collection = mock_db
+    documents = [{"_id": 1, "name": "test1"}, {"_id": 2, "name": "test2"}]
+    mock_collection.find.return_value = documents
+
+    result = get_data_from_collection(mock_db_instance, "test_collection")
+
+    mock_db_instance.__getitem__.assert_called_once_with("test_collection")
+    mock_collection.find.assert_called_once()
+    assert result is not None
+    assert isinstance(result[0], StringIO)
+    assert result[1] == "raw-files/test_collection.json"
+    expected_data = [{"name": "test1"}, {"name": "test2"}]
+    assert json.loads(result[0].getvalue()) == expected_data
+    mock_logger.info.assert_called_with("Extracting data from test_collection")
+
+
+def test_get_data_from_collection_no_documents(mock_db, mock_logger):
+    mock_db_instance, mock_collection = mock_db
+    mock_collection.find.return_value = []
+
+    result = get_data_from_collection(mock_db_instance, "test_collection")
+
+    mock_db_instance.__getitem__.assert_called_once_with("test_collection")
+    mock_collection.find.assert_called_once()
+    assert result is None
+    mock_logger.error.assert_called_with(f"No documents found in {mock_collection}")
+
+
+def test_get_data_from_collection_no_id_field(mock_db, mock_logger):
+    mock_db_instance, mock_collection = mock_db
+    documents = [{"name": "test1"}, {"name": "test2"}]
+    mock_collection.find.return_value = documents
+
+    result = get_data_from_collection(mock_db_instance, "test_collection")
+
+    mock_db_instance.__getitem__.assert_called_once_with("test_collection")
+    mock_collection.find.assert_called_once()
+    assert result is not None
+    assert isinstance(result[0], StringIO)
+    assert result[1] == "raw-files/test_collection.json"
+    expected_data = [{"name": "test1"}, {"name": "test2"}]
+    assert json.loads(result[0].getvalue()) == expected_data
+    mock_logger.info.assert_called_with("Extracting data from test_collection")