From e2ba1a7b8c228ab409f0488760aa05ac9717c89c Mon Sep 17 00:00:00 2001
From: Muhammad Nawzad <hama127n@gmail.com>
Date: Sun, 12 Nov 2023 10:29:19 +0300
Subject: [PATCH] Updates RelationshipSchemaGenerator to correctly generate
 schemas

---
 .../relationship_schema_generator.rb          | 48 ++++++++++++++-----
 .../relationship_schema_generator.rbs         |  2 +
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/lib/schemable/relationship_schema_generator.rb b/lib/schemable/relationship_schema_generator.rb
index f8b7e6e..609d75c 100644
--- a/lib/schemable/relationship_schema_generator.rb
+++ b/lib/schemable/relationship_schema_generator.rb
@@ -33,20 +33,9 @@ module Schemable
             }
           }
 
-          result = {
-            type: :object,
-            properties: {
-              data: {
-                type: :object,
-                properties: {
-                  id: { type: :string },
-                  type: { type: :string, default: definition[:definition].model_name }
-                }
-              }
-            }
-          }
+          result = relation_type == :belongs_to ? generate_schema(definition.model_name) : generate_schema(definition.model_name, collection: true)
 
-          result = non_expanded_data_properties if !expand || @relationships_to_exclude_from_expansion.include?(definition[:definition].model_name)
+          result = non_expanded_data_properties if !expand || @relationships_to_exclude_from_expansion.include?(definition.model_name)
 
           schema[:properties].merge!(relation => result)
         end
@@ -62,5 +51,38 @@ module Schemable
 
       schema
     end
+
+    def generate_schema(type_name, collection: false)
+      if collection
+        {
+          type: :object,
+          properties: {
+            data: {
+              type: :array,
+              items: {
+                type: :object,
+                properties: {
+                  id: { type: :string },
+                  type: { type: :string, default: type_name }
+                }
+              }
+            }
+          }
+        }
+      else
+        {
+          type: :object,
+          properties: {
+            data: {
+              type: :object,
+              properties: {
+                id: { type: :string },
+                type: { type: :string, default: type_name }
+              }
+            }
+          }
+        }
+      end
+    end
   end
 end
diff --git a/sig/schemable/relationship_schema_generator.rbs b/sig/schemable/relationship_schema_generator.rbs
index 9c1d7e5..a1c50e0 100644
--- a/sig/schemable/relationship_schema_generator.rbs
+++ b/sig/schemable/relationship_schema_generator.rbs
@@ -10,5 +10,7 @@ module Schemable
     def initialize: (Definition, ?relationships_to_exclude_from_expansion: Array[String], ?expand: bool) -> void
 
     def generate: -> (Hash[Symbol, any] | Array[Hash[Symbol, any]])
+
+    def generate_schema: (String, ?collection: bool) -> Hash[Symbol, any]
   end
 end