import pytest from unittest import mock from types import SimpleNamespace import sys # Import the functions to test from purge_ttl.py import purge_ttl def test_parse_args_list_single_item(): assert purge_ttl.parse_args_list("foo") == ["foo"] def test_parse_args_list_multiple_items(): assert purge_ttl.parse_args_list("[a,b,c]") == ["a", "b", "c"] def test_get_expiry_condition_now(): args = SimpleNamespace(expiry_mode="now") assert purge_ttl.get_expiry_condition(args) == 'expiry < CURRENT_TIMESTAMP()' def test_get_expiry_condition_midnight(): args = SimpleNamespace(expiry_mode="midnight") assert purge_ttl.get_expiry_condition(args) == 'expiry < TIMESTAMP_TRUNC(CURRENT_TIMESTAMP(), DAY, "UTC")' def test_get_expiry_condition_invalid(): args = SimpleNamespace(expiry_mode="invalid") with pytest.raises(Exception): purge_ttl.get_expiry_condition(args) def test_add_conditions_no_collections_no_prefix(): args = SimpleNamespace(collection_ids=[], uid_prefixes=None) query, params, types = purge_ttl.add_conditions(args, "SELECT * FROM foo WHERE 1=1", None) assert query == "SELECT * FROM foo WHERE 1=1" assert params == {} assert types == {} def test_add_conditions_with_collections_single(): args = SimpleNamespace(collection_ids=["123"]) query, params, types = purge_ttl.add_conditions(args, "SELECT * FROM foo WHERE 1=1", None) assert "collection_id = @collection_id" in query assert params["collection_id"] == "123" assert types["collection_id"] == purge_ttl.param_types.INT64 def test_add_conditions_with_collections_multiple(): args = SimpleNamespace(collection_ids=["1", "2"]) query, params, types = purge_ttl.add_conditions(args, "SELECT * FROM foo WHERE 1=1", None) assert "collection_id in" in query assert params["collection_id_0"] == "1" assert params["collection_id_1"] == "2" assert types["collection_id_0"] == purge_ttl.param_types.INT64 def test_add_conditions_with_prefix(): args = SimpleNamespace(collection_ids=[]) query, params, types = purge_ttl.add_conditions(args, "SELECT * FROM foo WHERE 1=1", "abc") assert "STARTS_WITH(fxa_uid, @prefix)" in query assert params["prefix"] == "abc" assert types["prefix"] == purge_ttl.param_types.STRING @mock.patch("purge_ttl.statsd") def test_deleter_dryrun(statsd_mock): database = mock.Mock() statsd_mock.timer.return_value.__enter__.return_value = None statsd_mock.timer.return_value.__exit__.return_value = None purge_ttl.deleter(database, "batches", "DELETE FROM batches", dryrun=True) database.execute_partitioned_dml.assert_not_called() @mock.patch("purge_ttl.statsd") def test_deleter_executes(statsd_mock): database = mock.Mock() statsd_mock.timer.return_value.__enter__.return_value = None statsd_mock.timer.return_value.__exit__.return_value = None database.execute_partitioned_dml.return_value = 42 purge_ttl.deleter(database, "batches", "DELETE FROM batches", dryrun=False) database.execute_partitioned_dml.assert_called_once() @mock.patch("purge_ttl.deleter") @mock.patch("purge_ttl.add_conditions") @mock.patch("purge_ttl.get_expiry_condition") @mock.patch("purge_ttl.client") def test_spanner_purge_both(client_mock, get_expiry_condition_mock, add_conditions_mock, deleter_mock): # Setup args = SimpleNamespace( instance_id="inst", database_id="db", expiry_mode="now", auto_split=None, uid_prefixes=None, mode="both", dryrun=True, collection_ids=[] ) instance = mock.Mock() database = mock.Mock() client_mock.instance.return_value = instance instance.database.return_value = database get_expiry_condition_mock.return_value = "expiry < CURRENT_TIMESTAMP()" add_conditions_mock.side_effect = [ ("batch_query", {"a": 1}, {"a": 2}), ("bso_query", {"b": 3}, {"b": 4}), ] purge_ttl.spanner_purge(args) assert deleter_mock.call_count == 2 deleter_mock.assert_called() deleter_mock.assert_any_call( database, name="batches", query="batch_query", params={"a": 1}, param_types={"a": 2}, prefix=None, dryrun=True, ) deleter_mock.assert_any_call( database, name="bso", query="bso_query", params={"b": 3}, param_types={"b": 4}, prefix=None, dryrun=True, ) @mock.patch("argparse.ArgumentParser.parse_args") def test_get_args_env_and_dsn(parse_args_mock): # Simulate args with DSN args = SimpleNamespace( instance_id="foo", database_id="bar", project_id="baz", sync_database_url="spanner://projects/proj/instances/inst/databases/db", collection_ids=[], uid_prefixes=[], auto_split=None, mode="both", expiry_mode="midnight", dryrun=False, ) parse_args_mock.return_value = args result = purge_ttl.get_args() assert result.project_id == "proj" assert result.instance_id == "inst" assert result.database_id == "db"