From e4d5710e74edb0e20b11190f16fcde79d8df97ed Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 18 Dec 2024 11:32:27 +0530
Subject: [PATCH] fix: retain system metadata when calling `flatten`

---
 src/systems/diffeqs/odesystem.jl              |  1 +
 .../discrete_system/discrete_system.jl        |  1 +
 src/systems/nonlinear/nonlinearsystem.jl      |  1 +
 .../optimization/optimizationsystem.jl        |  1 +
 test/components.jl                            | 36 +++++++++++++++++++
 5 files changed, 40 insertions(+)

diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
index 2b0bd8c8d7..34003f40c2 100644
--- a/src/systems/diffeqs/odesystem.jl
+++ b/src/systems/diffeqs/odesystem.jl
@@ -409,6 +409,7 @@ function flatten(sys::ODESystem, noeqs = false)
             initialization_eqs = initialization_equations(sys),
             is_dde = is_dde(sys),
             tstops = symbolic_tstops(sys),
+            metadata = get_metadata(sys),
             checks = false)
     end
 end
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
index 3e220998cb..01fca30235 100644
--- a/src/systems/discrete_system/discrete_system.jl
+++ b/src/systems/discrete_system/discrete_system.jl
@@ -227,6 +227,7 @@ function flatten(sys::DiscreteSystem, noeqs = false)
             defaults = defaults(sys),
             name = nameof(sys),
             description = description(sys),
+            metadata = get_metadata(sys),
             checks = false)
     end
 end
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
index 3cb68853aa..b2abac5184 100644
--- a/src/systems/nonlinear/nonlinearsystem.jl
+++ b/src/systems/nonlinear/nonlinearsystem.jl
@@ -859,6 +859,7 @@ function flatten(sys::NonlinearSystem, noeqs = false)
             defaults = defaults(sys),
             name = nameof(sys),
             description = description(sys),
+            metadata = get_metadata(sys),
             checks = false)
     end
 end
diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl
index 43e9294dd3..0398c892eb 100644
--- a/src/systems/optimization/optimizationsystem.jl
+++ b/src/systems/optimization/optimizationsystem.jl
@@ -184,6 +184,7 @@ function flatten(sys::OptimizationSystem)
         constraints = constraints(sys),
         defaults = defaults(sys),
         name = nameof(sys),
+        metadata = get_metadata(sys),
         checks = false
     )
 end
diff --git a/test/components.jl b/test/components.jl
index 298f9fceb9..8ac40f6fbb 100644
--- a/test/components.jl
+++ b/test/components.jl
@@ -370,3 +370,39 @@ end
     ss = structural_simplify(cbar)
     @test isequal(cbar.foo.x, ss.foo.x)
 end
+
+@testset "Issue#3275: Metadata retained on `complete`" begin
+    @variables x(t) y(t)
+    @testset "ODESystem" begin
+        @named inner = ODESystem(D(x) ~ x, t)
+        @named outer = ODESystem(D(y) ~ y, t; systems = [inner], metadata = "test")
+        @test ModelingToolkit.get_metadata(outer) == "test"
+        sys = complete(outer)
+        @test ModelingToolkit.get_metadata(sys) == "test"
+    end
+    @testset "NonlinearSystem" begin
+        @named inner = NonlinearSystem([0 ~ x^2 + 4x + 4], [x], [])
+        @named outer = NonlinearSystem(
+            [0 ~ x^3 - y^3], [x, y], []; systems = [inner], metadata = "test")
+        @test ModelingToolkit.get_metadata(outer) == "test"
+        sys = complete(outer)
+        @test ModelingToolkit.get_metadata(sys) == "test"
+    end
+    k = ShiftIndex(t)
+    @testset "DiscreteSystem" begin
+        @named inner = DiscreteSystem([x(k) ~ x(k - 1) + x(k - 2)], t, [x], [])
+        @named outer = DiscreteSystem([y(k) ~ y(k - 1) + y(k - 2)], t, [x, y],
+            []; systems = [inner], metadata = "test")
+        @test ModelingToolkit.get_metadata(outer) == "test"
+        sys = complete(outer)
+        @test ModelingToolkit.get_metadata(sys) == "test"
+    end
+    @testset "OptimizationSystem" begin
+        @named inner = OptimizationSystem(x^2 + y^2 - 3, [x, y], [])
+        @named outer = OptimizationSystem(
+            x^3 - y, [x, y], []; systems = [inner], metadata = "test")
+        @test ModelingToolkit.get_metadata(outer) == "test"
+        sys = complete(outer)
+        @test ModelingToolkit.get_metadata(sys) == "test"
+    end
+end