From c8991fd14cefb870daf36ee4e22fe71be98ea463 Mon Sep 17 00:00:00 2001
From: Faruk Alpay <32020561+farukalpay@users.noreply.github.com>
Date: Sun, 24 Aug 2025 06:35:28 +0200
Subject: [PATCH] test: add golden file for optimizer state

---
 tests/golden/psd_torch_state.json     | 25 +++++++++++++++++++
 tests/test_optimizer_serialization.py | 35 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 tests/golden/psd_torch_state.json
 create mode 100644 tests/test_optimizer_serialization.py

diff --git a/tests/golden/psd_torch_state.json b/tests/golden/psd_torch_state.json
new file mode 100644
index 0000000..44c20e9
--- /dev/null
+++ b/tests/golden/psd_torch_state.json
@@ -0,0 +1,25 @@
+{
+  "param_groups": [
+    {
+      "g_thres": 0.001,
+      "lr": 0.1,
+      "max_grad_norm": 1.0,
+      "params": [
+        0,
+        1
+      ],
+      "r": 0.001,
+      "t_thres": 10
+    }
+  ],
+  "state": {
+    "0": {
+      "t": 1,
+      "t_noise": -Infinity
+    },
+    "1": {
+      "t": 1,
+      "t_noise": -Infinity
+    }
+  }
+}
diff --git a/tests/test_optimizer_serialization.py b/tests/test_optimizer_serialization.py
new file mode 100644
index 0000000..63e8108
--- /dev/null
+++ b/tests/test_optimizer_serialization.py
@@ -0,0 +1,35 @@
+import json
+from pathlib import Path
+
+import torch
+
+from psd.framework_optimizers import PSDTorch
+
+
+def _compute_state_dict() -> dict:
+    model = torch.nn.Linear(1, 1)
+    model.weight.data.fill_(1.0)
+    model.bias.data.fill_(0.0)
+    opt = PSDTorch(model.parameters(), lr=0.1)
+
+    x = torch.tensor([[1.0]])
+    y = torch.tensor([[2.0]])
+    criterion = torch.nn.MSELoss()
+
+    def closure() -> torch.Tensor:
+        opt.zero_grad()
+        out = model(x)
+        loss = criterion(out, y)
+        loss.backward()
+        return loss
+
+    opt.step(closure)
+    return opt.state_dict()
+
+
+def test_psd_torch_state_dict_matches_golden() -> None:
+    state = _compute_state_dict()
+    generated = json.dumps(state, sort_keys=True, indent=2, allow_nan=True)
+    golden_path = Path(__file__).parent / "golden" / "psd_torch_state.json"
+    expected = golden_path.read_text().strip()
+    assert generated == expected