handling some more choice cases

open-traffic-generator · Mar 4, 2024 · 4218328 · 4218328
1 parent 55701b5
commit 4218328
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 37 deletions.
diff --git a/openapiart/openapiartgo.py b/openapiart/openapiartgo.py
@@ -3004,7 +3004,12 @@ def _write_default_method(self, new):
                         )
                     )
                 else:
-                    body += """if obj.obj.{name} == nil {{
+                    choice_cond = ""
+                    if field.name in choice_enum_map:
+                        choice_cond += (
+                            "&& choice == %s " % choice_enum_map[field.name]
+                        )
+                    body += """if obj.obj.{name} == nil {choice_check}{{
                         obj.Set{external_name}({value})
                     }}
                     """.format(
@@ -3015,6 +3020,7 @@ def _write_default_method(self, new):
                         value='"{0}"'.format(field.default)
                         if field.type == "string"
                         else field.default,
+                        choice_check=choice_cond,
                     )
             else:
                 if field.name in hasChoiceConfig:
@@ -3066,18 +3072,31 @@ def _write_default_method(self, new):
                     # signifies choice with no property
                     continue
 
-                enum_check_code = " && obj.obj.%s.Number() != 0 " % enum
-                choice_code += """
-                if obj.obj.{prop} != {val}{enum_check}{{
-                    choices_set += 1
-                    choice = {choice_val}
-                }}
-                """.format(
-                    prop=enum,
-                    val=value,
-                    choice_val=choice_enum_map[enum],
-                    enum_check=enum_check_code if field_type == "enum" else "",
-                )
+                if field_type.startswith("[]"):
+                    choice_code += """
+                    if len(obj.obj.{prop}) > 0{{
+                        choices_set += 1
+                        choice = {choice_val}
+                    }}
+                    """.format(
+                        prop=enum,
+                        choice_val=choice_enum_map[enum],
+                    )
+                else:
+                    enum_check_code = " && obj.obj.%s.Number() != 0 " % enum
+                    choice_code += """
+                    if obj.obj.{prop} != {val}{enum_check}{{
+                        choices_set += 1
+                        choice = {choice_val}
+                    }}
+                    """.format(
+                        prop=enum,
+                        val=value,
+                        choice_val=choice_enum_map[enum],
+                        enum_check=enum_check_code
+                        if field_type == "enum"
+                        else "",
+                    )
 
             # TODO: we need to throw error if more that one choice properties are set
             # choice_code += """

diff --git a/openapiart/tests/config/config.yaml b/openapiart/tests/config/config.yaml
@@ -319,6 +319,9 @@ components:
         choice_default:
           $ref: "#/components/schemas/ChoiceObject"
           x-field-uid: 56
+        choice_required_default:
+          $ref: "#/components/schemas/ChoiceRequiredAndDefault"
+          x-field-uid: 57
 
     WObject:
       required: [w_name]
@@ -738,3 +741,29 @@ components:
         ieee_802_3x:
           type: string
           x-field-uid: 5
+
+    ChoiceRequiredAndDefault:
+      type: object
+      required: [choice]
+      properties:
+        choice:
+          type: string
+          x-field-uid: 1
+          x-enum:
+            ipv4:
+              x-field-uid: 1
+            ipv6:
+              x-field-uid: 2
+        ipv4:
+          type: string
+          format: ipv4
+          default: "0.0.0.0"
+          x-field-uid: 2
+        ipv6:
+          description: |-
+            A list of ipv6
+          type: array
+          items:
+            type: string
+            format: ipv6
+          x-field-uid: 3
diff --git a/pkg/choice_test.go b/pkg/choice_test.go
@@ -421,41 +421,33 @@ func TestChoiceUnMarshall(t *testing.T) {
 
 func TestDefaultChoiceOverwrite(t *testing.T) {
 	config := openapiart.NewPrefixConfig()
-	fObj := config.F()
+	crd := config.ChoiceRequiredDefault()
+	crd.Ipv4()
 
 	// test default choice and values
-	fmt.Println(fObj)
-	assert.Equal(t, fObj.Choice(), openapiart.FObjectChoice.F_A)
-	assert.True(t, fObj.HasFA())
-	assert.Equal(t, fObj.FA(), "some string")
+	fmt.Println(crd)
+	assert.Equal(t, crd.Choice(), openapiart.ChoiceRequiredAndDefaultChoice.IPV4)
+	assert.True(t, crd.HasIpv4())
+	assert.Equal(t, crd.Ipv4(), "0.0.0.0")
 
 	// setting of other choices should work as usual
-	fObj.SetFB(5.67)
-
-	c, err := fObj.Marshal().ToYaml()
-	fmt.Println(c)
-	assert.Nil(t, err)
-
-	assert.Equal(t, fObj.Choice(), openapiart.FObjectChoice.F_B)
-	assert.True(t, fObj.HasFB())
-	assert.Equal(t, fObj.FB(), 5.67)
+	crd.SetIpv6([]string{"1::2"})
 
-	fObj.SetFA("str1")
-
-	c, err = fObj.Marshal().ToYaml()
+	c, err := crd.Marshal().ToYaml()
 	fmt.Println(c)
 	assert.Nil(t, err)
 
-	assert.Equal(t, fObj.Choice(), openapiart.FObjectChoice.F_A)
-	assert.True(t, fObj.HasFA())
-	assert.Equal(t, fObj.FA(), "str1")
+	assert.Equal(t, crd.Choice(), openapiart.ChoiceRequiredAndDefaultChoice.IPV6)
+	assert.True(t, len(crd.Ipv6()) > 0)
+	assert.Equal(t, crd.Ipv6()[0], "1::2")
 
-	// setting choice with no property
-	fObj.FC()
+	crd.SetIpv4("1.2.3.4")
 
-	c, err = fObj.Marshal().ToYaml()
+	c, err = crd.Marshal().ToYaml()
 	fmt.Println(c)
 	assert.Nil(t, err)
 
-	assert.Equal(t, fObj.Choice(), openapiart.FObjectChoice.F_C)
+	assert.Equal(t, crd.Choice(), openapiart.ChoiceRequiredAndDefaultChoice.IPV4)
+	assert.True(t, crd.HasIpv4())
+	assert.Equal(t, crd.Ipv4(), "1.2.3.4")
 }