From 526a0131f0208c399c3ead255b8bce7b9513c288 Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Tue, 3 Sep 2024 11:27:37 +1200
Subject: [PATCH] Add more tests for shapes and dual shapes (#3816)

---
 docs/src/developers/checklists.md | 14 +++++++++++
 test/test_constraint.jl           | 39 +++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/docs/src/developers/checklists.md b/docs/src/developers/checklists.md
index d9acdc39398..1693c623150 100644
--- a/docs/src/developers/checklists.md
+++ b/docs/src/developers/checklists.md
@@ -82,3 +82,17 @@ Use the following checklist when adding a new solver to the JuMP documentation.
 
  - [ ] Add package metadata to `docs/packages.toml`
 ````
+
+## Adding a new shape
+
+Use the following checklist when adding a new `AbstractShape`
+
+````
+## Basic
+
+ - [ ] Add a new subtype of `AbstractShape`
+ - [ ] Implement `vectorize(data, ::NewShape)::Vector`
+ - [ ] Implement `reshape_vector(vector, ::NewShape)`
+ - [ ] Implement `dual_shape`, or verify that the shape is self-dual
+ - [ ] Add the tests from https://github.com/jump-dev/JuMP.jl/pull/3816
+````
diff --git a/test/test_constraint.jl b/test/test_constraint.jl
index 2f183a36092..19c73209575 100644
--- a/test/test_constraint.jl
+++ b/test/test_constraint.jl
@@ -2118,4 +2118,43 @@ function test_issue_3812()
     return
 end
 
+function _test_shape(primal, dual, con::VectorConstraint)
+    vec_primal = vectorize(primal, con.shape)
+    vec_dual = vectorize(dual, dual_shape(con.shape))
+    @test primal == reshape_vector(vec_primal, con.shape)
+    @test dual == reshape_vector(vec_dual, dual_shape(con.shape))
+    moi_dot = MOI.Utilities.set_dot(vec_primal, vec_dual, con.set)
+    @test LinearAlgebra.dot(primal, dual) == moi_dot
+    return
+end
+
+function test_symmetric_adjoint_shape()
+    model = Model()
+    @variable(model, x[1:2, 1:2], Symmetric)
+    c = @constraint(model, x == LinearAlgebra.Symmetric([1 2; 2 3]))
+    _test_shape([1 2; 2 3], [1 1; 1 1], constraint_object(c))
+    return
+end
+
+function test_symmetric_shape()
+    model = Model()
+    @variable(model, x[1:2, 1:2], PSD)
+    c = @constraint(model, x == LinearAlgebra.Symmetric([1 2; 2 3]))
+    _test_shape([1 2; 2 3], [1 1; 1 1], constraint_object(c))
+    return
+end
+
+function test_hermitian_shape()
+    model = Model()
+    @variable(model, x[1:2, 1:2], Hermitian)
+    c = @constraint(model, x == LinearAlgebra.Symmetric([1 2; 2 3]))
+    _test_shape([1 2; 2 3], [2 4; 4 6], constraint_object(c))
+    model = Model()
+    @variable(model, x[1:2, 1:2], Hermitian)
+    primal = [1 2+4im; 2-4im 3]
+    d = @constraint(model, x == LinearAlgebra.Hermitian(primal))
+    _test_shape(primal, [2 4+8im; 4-8im 6], constraint_object(d))
+    return
+end
+
 end  # module