Repository: aidotse/LeakPro

Issue:
Unit test LiRA
# Issue
## Problem Description
We need a unit test for LiRA

## Expected Behavior
LiRA should run

## What Needs to be Done
Implement unit tests for all functions

## How Can It Be Tested or Reproduced
- Run unit test

Discussion:

Solution patch:
diff --git a/leakpro/attacks/utils/shadow_model_handler.py b/leakpro/attacks/utils/shadow_model_handler.py
index b67784c9..c528d8e7 100644
--- a/leakpro/attacks/utils/shadow_model_handler.py
+++ b/leakpro/attacks/utils/shadow_model_handler.py
@@ -30,6 +30,12 @@ def get_instance(*args, **kwargs):  # noqa: ANN003, ANN002, ANN202
             raise ValueError("Singleton already created with specific parameters.")
         return instances[cls]
 
+    def is_created() -> bool:
+        return cls in instances
+
+    # Attach a check method to the class
+    get_instance.is_created = is_created
+
     return get_instance
 
 @singleton


Test patch:
diff --git a/leakpro/tests/conftest.py b/leakpro/tests/conftest.py
index 94d9c842..21125ac5 100644
--- a/leakpro/tests/conftest.py
+++ b/leakpro/tests/conftest.py
@@ -1,17 +1,39 @@
 """Test for the image input handler."""
 import logging
 import os
+import pickle
+import shutil
+import copy
+import random 
+
 from typing import Generator
+from torch import randn_like, save
 
 import pytest
 from dotmap import DotMap
 
 from leakpro.tests.input_handler.image_utils import setup_image_test
 from leakpro.user_inputs.cifar10_input_handler import Cifar10InputHandler
-
-
-@pytest.fixture
-def image_handler() -> Generator[Cifar10InputHandler, None, None]:
+from leakpro.attacks.utils.shadow_model_handler import ShadowModelHandler
+from leakpro.tests.constants import STORAGE_PATH
+
+
+@pytest.fixture(scope="session")
+def manage_storage_directory():
+    """Fixture to create and remove the storage directory."""
+    
+    # Setup: Create the folder at the start of the test session
+    os.makedirs(STORAGE_PATH, exist_ok=True)
+    
+    # Yield control back to the test session
+    yield
+    
+    # Teardown: Remove the folder and its contents at the end of the session
+    if os.path.exists(STORAGE_PATH):
+        shutil.rmtree(STORAGE_PATH)
+
+@pytest.fixture()
+def image_handler(manage_storage_directory) -> Generator[Cifar10InputHandler, None, None]:
     """Fixture for the image input handler to be shared between many tests."""
 
     config = DotMap()
@@ -23,9 +45,5 @@ def image_handler() -> Generator[Cifar10InputHandler, None, None]:
 
     handler = Cifar10InputHandler(config, logger)
 
+    # Yield control back to the test session
     yield handler
-
-    # Clean up folder tests are done
-    os.remove(config.target.data_path) # remove data
-    os.remove(config.target.trained_model_path) # remove target model
-    os.remove(config.target.trained_model_metadata_path) # remove target metadata
diff --git a/leakpro/tests/constants.py b/leakpro/tests/constants.py
new file mode 100644
index 00000000..de373038
--- /dev/null
+++ b/leakpro/tests/constants.py
@@ -0,0 +1,34 @@
+from dotmap import DotMap
+
+STORAGE_PATH = "./leakpro/tests/tmp"
+
+# User input handler for images
+parameters = DotMap()
+parameters.epochs = 10
+parameters.batch_size = 64
+parameters.learning_rate = 0.001
+parameters.optimizer = "sgd"
+parameters.loss = "crossentropyloss"
+parameters.data_points = 130
+parameters.train_data_points = 20
+parameters.test_data_points = 20
+parameters.img_size = (3, 32, 32)
+parameters.num_classes = 13
+parameters.images_per_class = parameters.data_points // parameters.num_classes
+
+# Shadow model configuration for images
+shadow_model_config = DotMap()
+shadow_model_config.module_path = "./leakpro/tests/input_handler/image_utils.py"
+shadow_model_config.model_class = "ConvNet"
+shadow_model_config.storage_path = "./leakpro/tests/tmp/model_handler_output"
+shadow_model_config.batch_size = 32
+shadow_model_config.epochs = 1
+shadow_model_config.optimizer = {"name": "sgd", "lr": 0.001}
+shadow_model_config.loss = {"name": "crossentropyloss"}
+
+# Lira parameters
+lira_params = DotMap()
+lira_params.training_data_fraction = 0.1
+lira_params.num_shadow_models = 3
+lira_params.online = False
+lira_params.fixed_variance = True
\ No newline at end of file
diff --git a/leakpro/tests/input_handler/image_utils.py b/leakpro/tests/input_handler/image_utils.py
index 85ebbcde..61327d5d 100644
--- a/leakpro/tests/input_handler/image_utils.py
+++ b/leakpro/tests/input_handler/image_utils.py
@@ -12,22 +12,7 @@
 
 from leakpro.import_helper import Self
 
-STORAGE_PATH = "./leakpro/tests/input_handler"
-
-parameters = DotMap()
-parameters.epochs = 10
-parameters.batch_size = 64
-parameters.learning_rate = 0.001
-parameters.optimizer = "sgd"
-parameters.loss = "crossentropyloss"
-parameters.data_points = 130
-parameters.train_data_points = 20
-parameters.test_data_points = 20
-parameters.img_size = (3, 32, 32)
-parameters.num_classes = 13
-parameters.images_per_class = parameters.data_points // parameters.num_classes
-
-
+from leakpro.tests.constants import STORAGE_PATH, parameters
 
 class ConvNet(Module):
     """Convolutional Neural Network model."""
diff --git a/leakpro/tests/mia_attacks/attacks/test_lira.py b/leakpro/tests/mia_attacks/attacks/test_lira.py
new file mode 100644
index 00000000..68363141
--- /dev/null
+++ b/leakpro/tests/mia_attacks/attacks/test_lira.py
@@ -0,0 +1,134 @@
+
+from pytest import raises
+from math import isnan
+
+from leakpro.user_inputs.cifar10_input_handler import Cifar10InputHandler
+from leakpro.attacks.utils.shadow_model_handler import ShadowModelHandler
+from leakpro.attacks.mia_attacks.lira import AttackLiRA
+from leakpro.tests.constants import shadow_model_config, lira_params
+
+def test_lira_setup(image_handler:Cifar10InputHandler) -> None:
+    """Test the initialization of LiRA."""
+    
+    lira_obj = AttackLiRA(image_handler, lira_params)
+    
+    assert lira_obj is not None
+    assert lira_obj.target_model is not None
+    assert lira_obj.online == lira_params.online
+    assert lira_obj.num_shadow_models == lira_params.num_shadow_models
+    assert lira_obj.training_data_fraction == lira_params.training_data_fraction
+    assert lira_obj.memorization == False
+    
+    lira_params.num_shadow_models = -1
+    with raises(ValueError) as excinfo:
+        lira_obj = AttackLiRA(image_handler, lira_params)
+    assert str(excinfo.value) == "num_shadow_models must be between 1 and None"
+    
+    lira_params.num_shadow_models = 3
+    
+    description = lira_obj.description()
+    assert len(description) == 4
+
+def test_lira_prepare_online_attack(image_handler:Cifar10InputHandler) -> None:
+    lira_params.online = True
+    
+    image_handler.configs.shadow_model = shadow_model_config
+    lira_obj = AttackLiRA(image_handler, lira_params)
+    
+    if ShadowModelHandler.is_created() == False:
+        ShadowModelHandler(image_handler)
+        
+    lira_obj.prepare_attack()
+    
+    # ensure correct number of shadow models are read
+    assert len(lira_obj.shadow_models) == lira_params.num_shadow_models
+    # ensure the attack data indices correspond to the correct pool
+    assert sorted(lira_obj.attack_data_indices) == list(range(image_handler.population_size))
+    # memorization is tested in a different module
+    assert lira_obj.memorization == False
+    
+    # Check that the filtering of the attack data is correct (this is done after shadow models are created)
+    n_attack_points = len(lira_obj.in_members) + len(lira_obj.out_members)
+    assert n_attack_points > 0
+    assert lira_obj.shadow_models_logits.shape == (n_attack_points, lira_params.num_shadow_models)
+    assert lira_obj.target_logits.shape == (n_attack_points, )
+    
+def test_lira_prepare_offline_attack(image_handler:Cifar10InputHandler) -> None:
+    lira_params.online = False
+    
+    image_handler.configs.shadow_model = shadow_model_config
+    lira_obj = AttackLiRA(image_handler, lira_params)
+    
+    if ShadowModelHandler.is_created() == False:
+        ShadowModelHandler(image_handler)
+        
+    lira_obj.prepare_attack()
+    
+    # ensure correct number of shadow models are read
+    assert len(lira_obj.shadow_models) == lira_params.num_shadow_models
+    # ensure the attack data indices correspond to the correct pool
+    assert sorted(lira_obj.attack_data_indices) == sorted(set(range(image_handler.population_size)) - set(image_handler.test_indices) - set(image_handler.train_indices))
+    # memorization is tested in a different module
+    assert lira_obj.memorization == False
+    
+    # Check that the filtering of the attack data is correct (this is done after shadow models are created)
+    n_attack_points = len(lira_obj.in_members) + len(lira_obj.out_members)
+    assert n_attack_points > 0
+    assert lira_obj.shadow_models_logits.shape == (n_attack_points, lira_params.num_shadow_models)
+    assert lira_obj.target_logits.shape == (n_attack_points, )
+
+
+def test_lira_online_attack(image_handler:Cifar10InputHandler):
+    # Set up for testing
+    lira_params.online = True
+    image_handler.configs.shadow_model = shadow_model_config
+    lira_obj = AttackLiRA(image_handler, lira_params)
+    if ShadowModelHandler.is_created() == False:
+        ShadowModelHandler(image_handler)
+    lira_obj.prepare_attack()
+    
+    # Test standard deviation calculation
+    std_fixed = lira_obj.get_std(lira_obj.shadow_models_logits.flatten(),
+                           lira_obj.in_indices_masks.flatten(),
+                           True,
+                           "fixed")
+    lira_obj.fixed_in_std = std_fixed
+    lira_obj.fixed_out_std = std_fixed
+
+    std_carlini = lira_obj.get_std(lira_obj.shadow_models_logits.flatten(),
+                           lira_obj.in_indices_masks.flatten(),
+                           True,
+                           "carlini")
+    
+    std_individual = lira_obj.get_std(lira_obj.shadow_models_logits.flatten(),
+                           lira_obj.in_indices_masks.flatten(),
+                           True,
+                           "individual_carlini")
+    assert std_fixed == std_carlini
+    assert std_fixed == std_individual
+    
+    # Test attack
+    lira_obj.run_attack()
+    assert lira_obj.fixed_in_std != lira_obj.fixed_out_std
+    n_attack_points = len(lira_obj.in_members) + len(lira_obj.out_members)
+    assert len(lira_obj.in_member_signals)+len(lira_obj.out_member_signals) == n_attack_points
+    assert any(isnan(x) for x in lira_obj.in_member_signals) == False
+    assert any(isnan(x) for x in lira_obj.out_member_signals) == False
+    
+def test_lira_online_attack(image_handler:Cifar10InputHandler):
+    # Set up for testing
+    lira_params.online = False
+    image_handler.configs.shadow_model = shadow_model_config
+    lira_obj = AttackLiRA(image_handler, lira_params)
+    if ShadowModelHandler.is_created() == False:
+        ShadowModelHandler(image_handler)
+    lira_obj.prepare_attack()
+    lira_obj.fix_var_threshold = 0.0
+    
+    # Test attack
+    lira_obj.run_attack()
+    assert lira_obj.fixed_in_std != lira_obj.fixed_out_std
+    n_attack_points = len(lira_obj.in_members) + len(lira_obj.out_members)
+    assert len(lira_obj.in_member_signals)+len(lira_obj.out_member_signals) == n_attack_points
+    assert any(isnan(x) for x in lira_obj.in_member_signals) == False
+    assert any(isnan(x) for x in lira_obj.out_member_signals) == False
\ No newline at end of file
diff --git a/leakpro/tests/mia_attacks/utils/test_shadow_model_handler.py b/leakpro/tests/mia_attacks/utils/test_shadow_model_handler.py
index 007bf83d..c04550ae 100644
--- a/leakpro/tests/mia_attacks/utils/test_shadow_model_handler.py
+++ b/leakpro/tests/mia_attacks/utils/test_shadow_model_handler.py
@@ -1,42 +1,32 @@
 """Test the shadow model handler module."""
 import os
-import shutil
-from dotmap import DotMap
 import numpy as np
 
 from pytest import raises
 from leakpro.attacks.utils.shadow_model_handler import ShadowModelHandler
 from leakpro.user_inputs.cifar10_input_handler import Cifar10InputHandler
-
-def _setup_shadow_test() -> DotMap:
-    """Setup the shadow model test."""
-    shadow_model_config = DotMap()
-    shadow_model_config.module_path = "./leakpro/tests/input_handler/image_utils.py"
-    shadow_model_config.model_class = "ConvNet"
-    shadow_model_config.storage_path = "./leakpro/tests/model_handlers/model_handler_output"
-    shadow_model_config.batch_size = 32
-    shadow_model_config.epochs = 1
-    shadow_model_config.optimizer = {"name": "sgd", "lr": 0.001}
-    shadow_model_config.loss = {"name": "crossentropyloss"}
-    return shadow_model_config
-
+from leakpro.tests.constants import shadow_model_config
 
 def test_shadow_model_handler_singleton(image_handler:Cifar10InputHandler) -> None:
     """Test that only one instance gets created."""
-    shadow_model_config = _setup_shadow_test()
+    
     image_handler.configs.shadow_model = shadow_model_config
-    sm = ShadowModelHandler(image_handler)
+    if ShadowModelHandler.is_created() == False:
+        sm = ShadowModelHandler(image_handler)
+        assert ShadowModelHandler.is_created() == True
     
     with raises(ValueError) as excinfo:
         ShadowModelHandler(image_handler)
     assert str(excinfo.value) == "Singleton already created with specific parameters."
 
 def test_shadow_model_creation_and_loading(image_handler:Cifar10InputHandler) -> None:
-    shadow_model_config = _setup_shadow_test()
     image_handler.configs.shadow_model = shadow_model_config
     
     # Test initialization
-    sm = ShadowModelHandler()
+    if ShadowModelHandler.is_created() == False:
+        sm = ShadowModelHandler(image_handler)
+    else:
+        sm = ShadowModelHandler()
     
     assert sm.batch_size == shadow_model_config.batch_size
     assert sm.epochs == shadow_model_config.epochs
@@ -50,38 +40,39 @@ def test_shadow_model_creation_and_loading(image_handler:Cifar10InputHandler) ->
     training_fraction = 0.5
     online = False
     
-    entries = os.listdir(sm.storage_path)
-    assert len(entries) == 0
-    
-    sm.create_shadow_models(n_models, image_handler.test_indices, training_fraction, online)
+    entries_start = os.listdir(sm.storage_path)
+    n_entries_start = len(entries_start)
     
+    indx = sm.create_shadow_models(n_models, image_handler.test_indices, training_fraction, online)[0]
     entries = os.listdir(sm.storage_path)
-    assert len(entries) == 2*n_models
-    assert "metadata_0.pkl" in entries
-    assert "shadow_model_0.pkl" in entries
+    n_entries_phase1 = len(entries)
+    assert n_entries_phase1 - n_entries_start == 2*n_models
+    assert f"metadata_{indx}.pkl" in entries
+    assert f"shadow_model_{indx}.pkl" in entries
     
-    sm.create_shadow_models(n_models, image_handler.test_indices, training_fraction, ~online)
-    trained_models = 2
+
+    indx2 = sm.create_shadow_models(n_models, image_handler.test_indices, training_fraction, ~online)[0]
     entries = os.listdir(sm.storage_path)
-    assert len(entries) == 2*trained_models
-    assert "metadata_0.pkl" in entries
-    assert "shadow_model_0.pkl" in entries
-    assert "metadata_1.pkl" in entries
-    assert "shadow_model_1.pkl" in entries
+    assert len(entries) - n_entries_phase1 == 2*n_models
+    assert f"metadata_{indx}.pkl" in entries
+    assert f"shadow_model_{indx}.pkl" in entries
+    assert f"metadata_{indx2}.pkl" in entries
+    assert f"shadow_model_{indx2}.pkl" in entries
     
     # Test loading
-    meta_0 = sm._load_metadata(sm.storage_path + "/metadata_0.pkl")
-    meta_1 = sm._load_metadata(sm.storage_path + "/metadata_1.pkl")
+    meta_0 = sm._load_metadata(sm.storage_path + f"/metadata_{indx}.pkl")
+    meta_1 = sm._load_metadata(sm.storage_path + f"/metadata_{indx2}.pkl")
     assert meta_0["online"] == online
     assert meta_1["online"] == ~online
     assert meta_1["num_train"] == meta_0["num_train"]
     
-    shadow_model_indices = [0,1]
+    shadow_model_indices = [indx, indx2]
     models, indices = sm.get_shadow_models(shadow_model_indices)
     for model in models:
         assert model.model_obj.__class__.__name__ == "ConvNet"
     
     # Test index mask
+    # check what test data is included in the training data of the shadow models
     mask = sm.get_in_indices_mask(shadow_model_indices, np.array(image_handler.test_indices))
     
     true_mask_0 = np.array([True if item in meta_0["train_indices"] else False for item in image_handler.test_indices])
@@ -89,8 +80,3 @@ def test_shadow_model_creation_and_loading(image_handler:Cifar10InputHandler) ->
     
     true_mask_1 = np.array([True if item in meta_1["train_indices"] else False for item in image_handler.test_indices])
     assert np.array_equal(mask[:, 1], true_mask_1)
-    
-    # remove created files
-    shutil.rmtree(sm.storage_path)
-    del sm
-    
\ No newline at end of file