From 4f89f08b7474469434ea38e3a1bcc5d7f5643d3e Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Wed, 30 Oct 2024 12:05:24 +0530
Subject: [PATCH] fix: fix `complete` not properly expanding systems

---
 src/systems/abstractsystem.jl |  4 ++--
 test/odesystem.jl             | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
index de389ab3b8..c53c900d27 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -933,8 +933,8 @@ function complete(sys::AbstractSystem; split = true, flatten = true)
         @set! sys.ps = unique!(vcat(get_ps(sys), collect(newparams)))
     end
     if flatten
-        if (eqs = equations(sys)) isa Vector &&
-           any(eq -> eq isa Equation && isconnection(eq.lhs), eqs)
+        eqs = equations(sys)
+        if eqs isa AbstractArray && eltype(eqs) <: Equation
             newsys = expand_connections(sys)
         else
             newsys = sys
diff --git a/test/odesystem.jl b/test/odesystem.jl
index 9286fb3c01..dfc755da98 100644
--- a/test/odesystem.jl
+++ b/test/odesystem.jl
@@ -1467,3 +1467,31 @@ end
     obsfn(buf, prob.u0, prob.p, 0.0)
     @test buf ≈ [1.0, 1.0, 2.0]
 end
+
+@testset "`complete` expands connections" begin
+    using ModelingToolkitStandardLibrary.Electrical
+    @mtkmodel RC begin
+        @parameters begin
+            R = 1.0
+            C = 1.0
+            V = 1.0
+        end
+        @components begin
+            resistor = Resistor(R = R)
+            capacitor = Capacitor(C = C, v = 0.0)
+            source = Voltage()
+            constant = Constant(k = V)
+            ground = Ground()
+        end
+        @equations begin
+            connect(constant.output, source.V)
+            connect(source.p, resistor.p)
+            connect(resistor.n, capacitor.p)
+            connect(capacitor.n, source.n, ground.g)
+        end
+    end
+    @named sys = RC()
+    total_eqs = length(equations(expand_connections(sys)))
+    sys2 = complete(sys)
+    @test length(equations(sys2)) == total_eqs
+end