test_lira_attack.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pyre-strict


import unittest

import numpy as np
import pandas as pd
from privacy_guard.analysis.mia.aggregate_analysis_input import (
    AggregateAnalysisInput,
    AggregationType,
)
from privacy_guard.attacks.lira_attack import LiraAttack


class TestLiraAttack(unittest.TestCase):
    def setUp(self) -> None:
        self.df_train_merge = {
            "user_id": {
                "0": 100001,
                "1": 100002,
                "2": 100003,
                "3": 100004,
                "4": 100005,
            },
            "sample_id": {
                "0": 101,
                "1": 102,
                "2": 103,
                "3": 104,
                "4": 105,
            },
            "timestamp": {
                "0": 1700000000,
                "1": 1700000001,
                "2": 1700000002,
                "3": 1700000003,
                "4": 1700000004,
            },
            "hash_id": {
                "0": 1.0,
                "1": 20.0,
                "2": 300.0,
                "3": 4000.0,
                "4": 50000.0,
            },
            "score_orig": {"0": 0.2, "1": 0.4, "2": 0.5, "3": 0.67, "4": 0.99},
            "score_mean": {"0": 0.2, "1": 0.4, "2": 0.5, "3": 0.67, "4": 0.99},
            "score_std": {"0": 0.2, "1": 0.2, "2": 0.5, "3": 0.67, "4": 0.2},
            # Additional columns for online attack tests
            "score_mean_in": {"0": 0.25, "1": 0.45, "2": 0.55, "3": 0.7, "4": 0.95},
            "score_mean_out": {"0": 0.15, "1": 0.35, "2": 0.45, "3": 0.6, "4": 0.9},
            "score_std_in": {"0": 0.1, "1": 0.15, "2": 0.2, "3": 0.25, "4": 0.3},
            "score_std_out": {"0": 0.05, "1": 0.1, "2": 0.15, "3": 0.2, "4": 0.25},
        }
        self.df_train_merge = pd.DataFrame.from_dict(self.df_train_merge)

        self.user_id_key = "user_id"

        self.lira_attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=True,
            user_id_key=self.user_id_key,
        )

        self.lira_attack_no_fixed_variance = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=False,
            user_id_key=self.user_id_key,
        )

        super().setUp()

    def test_run_attack_presto_mocked(self) -> None:
        analysis_input = self.lira_attack.run_attack()
        self.assertTrue(isinstance(analysis_input, AggregateAnalysisInput))

        self.assertTrue(analysis_input is not None)
        self.assertEqual(len(analysis_input.df_train_merge), 5)
        self.assertEqual(len(analysis_input.df_test_merge), 5)

    def test_run_attack_presto_mocked_no_fixed_variance(self) -> None:
        analysis_input = self.lira_attack_no_fixed_variance.run_attack()
        self.assertTrue(isinstance(analysis_input, AggregateAnalysisInput))

        self.assertTrue(analysis_input is not None)
        self.assertEqual(len(analysis_input.df_train_merge), 5)
        self.assertEqual(len(analysis_input.df_test_merge), 5)

    def test_get_std_dev_global(self) -> None:
        """Test _get_std_dev with std_dev_type='global'"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="global",
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected standard deviation of all score_orig values
        expected_std = pd.concat(
            [
                self.df_train_merge.score_orig,
                self.df_train_merge.score_orig,  # df_test_merge is same as df_train_merge in this test
            ]
        ).std()

        self.assertAlmostEqual(float(std_in), float(expected_std))
        self.assertAlmostEqual(float(std_out), float(expected_std))
        self.assertEqual(
            std_in, std_out
        )  # For global, std_in and std_out should be equal

    def test_get_std_dev_shadows_in_offline(self) -> None:
        """Test _get_std_dev with std_dev_type='shadows_in' and online_attack=False"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="shadows_in",
            online_attack=False,
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected mean of all score_std values
        expected_std = pd.concat(
            [
                self.df_train_merge.score_std,
                self.df_train_merge.score_std,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        self.assertAlmostEqual(float(std_in), float(expected_std))
        self.assertAlmostEqual(float(std_out), float(expected_std))
        self.assertEqual(
            std_in, std_out
        )  # For shadows_in offline, std_in and std_out should be equal

    def test_get_std_dev_shadows_in_online(self) -> None:
        """Test _get_std_dev with std_dev_type='shadows_in' and online_attack=True"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="shadows_in",
            online_attack=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected mean of all score_std_in values
        expected_std = pd.concat(
            [
                self.df_train_merge.score_std_in,
                self.df_train_merge.score_std_in,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        self.assertAlmostEqual(float(std_in), float(expected_std))
        self.assertAlmostEqual(float(std_out), float(expected_std))
        self.assertEqual(
            std_in, std_out
        )  # For shadows_in online, std_in and std_out should be equal

    def test_get_std_dev_shadows_out_offline(self) -> None:
        """Test _get_std_dev with std_dev_type='shadows_out' and online_attack=False"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="shadows_out",
            online_attack=False,
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected mean of all score_std values
        expected_std = pd.concat(
            [
                self.df_train_merge.score_std,
                self.df_train_merge.score_std,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        self.assertAlmostEqual(float(std_in), float(expected_std))
        self.assertAlmostEqual(float(std_out), float(expected_std))
        self.assertEqual(
            std_in, std_out
        )  # For shadows_out offline, std_in and std_out should be equal

    def test_get_std_dev_shadows_out_online(self) -> None:
        """Test _get_std_dev with std_dev_type='shadows_out' and online_attack=True"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="shadows_out",
            online_attack=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected mean of all score_std_out values
        expected_std = pd.concat(
            [
                self.df_train_merge.score_std_out,
                self.df_train_merge.score_std_out,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        self.assertAlmostEqual(float(std_in), float(expected_std))
        self.assertAlmostEqual(float(std_out), float(expected_std))
        self.assertEqual(
            std_in, std_out
        )  # For shadows_out online, std_in and std_out should be equal

    def test_get_std_dev_mix(self) -> None:
        """Test _get_std_dev with std_dev_type='mix' and online_attack=True"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="mix",
            online_attack=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        std_in, std_out = attack._get_std_dev()

        # Verify
        # Calculate expected mean of all score_std_in values for std_in
        expected_std_in = pd.concat(
            [
                self.df_train_merge.score_std_in,
                self.df_train_merge.score_std_in,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        # Calculate expected mean of all score_std_out values for std_out
        expected_std_out = pd.concat(
            [
                self.df_train_merge.score_std_out,
                self.df_train_merge.score_std_out,  # df_test_merge is same as df_train_merge in this test
            ]
        ).mean()

        self.assertAlmostEqual(float(std_in), float(expected_std_in))
        self.assertAlmostEqual(float(std_out), float(expected_std_out))
        self.assertNotEqual(
            std_in, std_out
        )  # For mix, std_in and std_out should be different

    def test_get_std_dev_mix_offline_error(self) -> None:
        """Test _get_std_dev with std_dev_type='mix' and online_attack=False raises ValueError"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="mix",
            online_attack=False,
            user_id_key=self.user_id_key,
        )

        # Execute and Verify
        with self.assertRaises(ValueError) as context:
            attack._get_std_dev()

        self.assertIn(
            "mix std dev type is only supported for online attacks",
            str(context.exception),
        )

    def test_get_std_dev_invalid_type(self) -> None:
        """Test _get_std_dev with invalid std_dev_type raises ValueError"""
        # Setup
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            std_dev_type="invalid_type",
            user_id_key=self.user_id_key,
        )

        # Execute and Verify
        with self.assertRaises(ValueError) as context:
            attack._get_std_dev()

        self.assertIn("is not a valid std_dev type", str(context.exception))

    def test_run_attack_drops_nan_rows_in_train(self) -> None:
        """Test that run_attack drops rows with NaN values in df_train_merge after logpdf computation."""
        # Setup: create training data with NaN in score_orig so logpdf produces NaN
        df_train_with_nan = self.df_train_merge.copy()
        df_train_with_nan.loc["0", "score_orig"] = np.nan
        df_train_with_nan.loc["2", "score_orig"] = np.nan

        attack = LiraAttack(
            df_train_merge=df_train_with_nan,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=True,
            user_id_key=self.user_id_key,
            online_attack=True,
        )

        # Execute
        analysis_input = attack.run_attack()

        # Assert: 2 NaN rows dropped from train, test unchanged
        self.assertIsInstance(analysis_input, AggregateAnalysisInput)
        assert isinstance(analysis_input, AggregateAnalysisInput)
        self.assertEqual(len(analysis_input.df_train_merge), 3)
        self.assertEqual(len(analysis_input.df_test_merge), 5)

    def test_run_attack_drops_nan_rows_in_test(self) -> None:
        """Test that run_attack drops rows with NaN values in df_test_merge after logpdf computation."""
        # Setup: create test data with NaN in score_orig so logpdf produces NaN
        df_test_with_nan = self.df_train_merge.copy()
        df_test_with_nan.loc["1", "score_orig"] = np.nan

        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=df_test_with_nan,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        analysis_input = attack.run_attack()

        # Assert: train unchanged, 1 NaN row dropped from test
        self.assertIsInstance(analysis_input, AggregateAnalysisInput)
        assert isinstance(analysis_input, AggregateAnalysisInput)
        self.assertEqual(len(analysis_input.df_train_merge), 5)
        self.assertEqual(len(analysis_input.df_test_merge), 4)

    def test_run_attack_drops_nan_rows_online_attack(self) -> None:
        """Test that run_attack drops NaN rows for online attack mode."""
        # Setup: create data with NaN in score_mean_in to produce NaN in logpdf
        df_train_with_nan = self.df_train_merge.copy()
        df_train_with_nan.loc["0", "score_mean_in"] = np.nan
        df_train_with_nan.loc["3", "score_mean_out"] = np.nan

        attack = LiraAttack(
            df_train_merge=df_train_with_nan,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=True,
            online_attack=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        analysis_input = attack.run_attack()

        # Assert: 2 NaN rows dropped from train, test unchanged
        self.assertIsInstance(analysis_input, AggregateAnalysisInput)
        assert isinstance(analysis_input, AggregateAnalysisInput)
        self.assertEqual(len(analysis_input.df_train_merge), 3)
        self.assertEqual(len(analysis_input.df_test_merge), 5)

    def test_run_attack_no_nan_preserves_all_rows(self) -> None:
        """Test that run_attack preserves all rows when no NaN values are present."""
        # Setup: use clean data (no NaN)
        attack = LiraAttack(
            df_train_merge=self.df_train_merge,
            df_test_merge=self.df_train_merge,
            row_aggregation=AggregationType.MAX,
            use_fixed_variance=True,
            user_id_key=self.user_id_key,
        )

        # Execute
        analysis_input = attack.run_attack()

        # Assert: all rows preserved
        self.assertIsInstance(analysis_input, AggregateAnalysisInput)
        assert isinstance(analysis_input, AggregateAnalysisInput)
        self.assertEqual(len(analysis_input.df_train_merge), 5)
        self.assertEqual(len(analysis_input.df_test_merge), 5)