diff --git a/src/Umbraco.TestData/SegmentTestController.cs b/src/Umbraco.TestData/SegmentTestController.cs
index 33badbbb55..650820760e 100644
--- a/src/Umbraco.TestData/SegmentTestController.cs
+++ b/src/Umbraco.TestData/SegmentTestController.cs
@@ -30,9 +30,8 @@ namespace Umbraco.TestData
             if (ct.Variations.VariesBySegment())
                 return Content($"The document type {alias} already allows segments, nothing has been changed");
 
-            ct.Variations = ct.Variations.SetFlag(ContentVariation.Segment);
-            
-            propType.Variations = propType.Variations.SetFlag(ContentVariation.Segment);
+            ct.SetVariesBy(ContentVariation.Segment);
+            propType.SetVariesBy(ContentVariation.Segment);
 
             Services.ContentTypeService.Save(ct);
             return Content($"The document type {alias} and property type {propertyTypeAlias} now allows segments");
@@ -50,7 +49,7 @@ namespace Umbraco.TestData
             if (!ct.VariesBySegment())
                 return Content($"The document type {alias} does not allow segments, nothing has been changed");
 
-            ct.Variations = ct.Variations.UnsetFlag(ContentVariation.Segment);
+            ct.SetVariesBy(ContentVariation.Segment, false);
 
             Services.ContentTypeService.Save(ct);
             return Content($"The document type {alias} no longer allows segments");