mobility-feed-api/functions-python/dataset_service/tests/test_dataset_service.py at 1aa0298735ce1fea198f1cb467233a0d102fbbfe · MobilityData/mobility-feed-api · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import unittest
from datetime import datetime
from unittest.mock import patch, MagicMock

from dataset_service_commons import DatasetTrace, Status, BatchExecution
from main import DatasetTraceService, BatchExecutionService


class TestDatasetService(unittest.TestCase):
    @patch("google.cloud.datastore.Client")
    def test_save_dataset_trace(self, mock_datastore_client):
        service = DatasetTraceService(mock_datastore_client)
        dataset_trace = DatasetTrace(
            stable_id="123", status=Status.PUBLISHED, timestamp=datetime.now()
        )
        service.save(dataset_trace)
        mock_datastore_client.put.assert_called_once()

    @patch("google.cloud.datastore.Client")
    def test_validate_and_save_exception(self, mock_datastore_client):
        service = DatasetTraceService(mock_datastore_client)
        dataset_trace = DatasetTrace(
            stable_id="123", status=Status.PUBLISHED, timestamp=datetime.now()
        )
        with self.assertRaises(ValueError):
            service.validate_and_save(dataset_trace, 1)

    @patch("google.cloud.datastore.Client")
    def test_validate_and_save(self, mock_datastore_client):
        service = DatasetTraceService(mock_datastore_client)
        dataset_trace = DatasetTrace(
            stable_id="123",
            execution_id="123",
            status=Status.PUBLISHED,
            timestamp=datetime.now(),
        )
        service.validate_and_save(dataset_trace, 1)
        mock_datastore_client.put.assert_called_once()

    @patch("google.cloud.datastore.Client")
    def test_get_dataset_trace_by_id(self, mock_datastore_client):
        mock_datastore_client.get.return_value = {
            "trace_id": "123",
            "stable_id": "123",
            "status": "PUBLISHED",
            "timestamp": datetime.now(),
        }
        service = DatasetTraceService(mock_datastore_client)
        trace = service.get_by_id("123")
        self.assertEqual(trace.stable_id, "123")

    @patch("google.cloud.datastore.Client")
    def test_get_dataset_trace_by_id_not_exist(self, mock_datastore_client):
        mock_datastore_client.get.return_value = None
        service = DatasetTraceService(mock_datastore_client)
        trace = service.get_by_id("non existent trace")
        self.assertListEqual(trace, [])

    @patch("google.cloud.datastore.Client")
    def test_save_batch_execution(self, mock_datastore_client):
        mock_datastore_client.return_value.key = MagicMock()
        mock_datastore_client.return_value.Entity = MagicMock()
        service = BatchExecutionService()
        service.client = mock_datastore_client
        execution = BatchExecution(
            execution_id="123", timestamp=datetime.now(), feeds_total=10
        )
        service.save(execution)
        mock_datastore_client.key.assert_called_once_with("batch_execution", "123")

    @patch("google.cloud.datastore.Client")
    @patch("main.DatasetTraceService._entity_to_dataset_trace")
    def test_get_by_execution_and_stable_ids(
        self, mock_entity_to_dataset_trace, mock_datastore_client
    ):
        mock_query = MagicMock()
        mock_datastore_client.query.return_value = mock_query

        mock_entity = MagicMock(
            spec=dict,
            **{
                "get.side_effect": lambda key: {
                    "trace_id": "123",
                    "stable_id": "stable_123",
                    "status": "PUBLISHED",
                    "timestamp": datetime.now(),
                }.get(key, None)
            }
        )
        mock_query.fetch.return_value = [mock_entity]
        mock_entity_to_dataset_trace.side_effect = lambda entity: DatasetTrace(
            stable_id=entity.get("stable_id"),
            status=Status(entity.get("status")),
            timestamp=entity.get("timestamp"),
            trace_id=entity.get("trace_id"),
        )
        service = DatasetTraceService(client=mock_datastore_client)

        results = service.get_by_execution_and_stable_ids("exec_id", "stable_123")

        mock_datastore_client.query.assert_called_once_with(kind="dataset_trace")
        mock_query.add_filter.assert_any_call("execution_id", "=", "exec_id")
        mock_query.add_filter.assert_any_call("stable_id", "=", "stable_123")
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].stable_id, "stable_123")