Upload folder using huggingface_hub

__init__.py

@@ -1,9 +1 @@
-# AnonySpark: Lightweight PySpark data anonymization
-from .masking import (
-    mask_email, mask_name, mask_date,
-    mask_ssn, mask_itin, mask_phone,
-    mask_email_udf, mask_name_udf, mask_date_udf,
-    mask_ssn_udf, mask_itin_udf, mask_phone_udf
-)
-
-from .utils import apply_masking
+# Init for tests

test_masking.py

@@ -0,0 +1,32 @@
+from anonyspark.masking import (
+    mask_email, mask_name, mask_date,
+    mask_ssn, mask_itin, mask_phone
+)
+
+def test_mask_email():
+    assert mask_email("[email protected]") == "***@example.com"
+    assert mask_email("") is None
+    assert mask_email(None) is None
+
+def test_mask_name():
+    assert mask_name("John") == "J***"
+    assert mask_name("") is None
+    assert mask_name(None) is None
+
+def test_mask_date():
+    assert mask_date("1991-08-14") == "***-**-14"
+    assert mask_date("invalid") is None
+    assert mask_date(None) is None
+
+def test_mask_ssn():
+    assert mask_ssn("123-45-6789") == "***-**-6789"
+    assert mask_ssn("invalid") is None
+
+def test_mask_itin():
+    assert mask_itin("912-73-1234") == "***-**-1234"
+    assert mask_itin("123-45-6789") is None
+
+def test_mask_phone():
+    assert mask_phone("123-456-7890") == "***-***-7890"
+    assert mask_phone("(123) 456-7890") == "***-***-7890"
+    assert mask_phone("invalid") is None

test_schema_masking.py

@@ -0,0 +1,40 @@
+# tests/test_schema_masking.py
+
+import sys
+import os
+
+sys.path.append("/content/anonyspark")
+
+from pyspark.sql import SparkSession
+from anonyspark.utils import apply_masking
+
+def test_schema_masking():
+    spark = SparkSession.builder.master("local[1]").appName("Test").getOrCreate()
+
+    df = spark.createDataFrame([{
+        "email": "[email protected]",
+        "name": "John",
+        "dob": "1991-08-14",
+        "ssn": "123-45-6789",
+        "itin": "912-73-1234",
+        "phone": "123-456-7890"
+    }])
+
+    schema = {
+        "email": "email",
+        "name": "name",
+        "dob": "dob",
+        "ssn": "ssn",
+        "itin": "itin",
+        "phone": "phone"
+    }
+
+    masked_df = apply_masking(df, schema)
+    result = masked_df.collect()[0].asDict()
+
+    assert result["masked_email"] == "***@example.com"
+    assert result["masked_name"] == "J***"
+    assert result["masked_dob"] == "***-**-14"
+    assert result["masked_ssn"] == "***-**-6789"
+    assert result["masked_itin"] == "***-**-1234"
+    assert result["masked_phone"] == "***-***-7890"