1313from pals_schema .ThickElement import ThickElement
1414from pals_schema .DriftElement import DriftElement
1515from pals_schema .QuadrupoleElement import QuadrupoleElement
16- from pals_schema .Line import Line
16+ from pals_schema .BeamLine import BeamLine
1717
1818
1919def test_BaseElement ():
@@ -98,23 +98,23 @@ def test_QuadrupoleElement():
9898 assert element .MagneticMultipoleP .Bn2 == element_magnetic_multipole_Bn2
9999 assert element .MagneticMultipoleP .Bs2 == element_magnetic_multipole_Bs2
100100 assert element .MagneticMultipoleP .tilt2 == element_magnetic_multipole_tilt2
101- # Serialize the Line object to YAML
101+ # Serialize the BeamLine object to YAML
102102 yaml_data = yaml .dump (element .model_dump (), default_flow_style = False )
103103 print (f"\n { yaml_data } " )
104104
105105
106- def test_Line ():
106+ def test_BeamLine ():
107107 # Create first line with one base element
108108 element1 = BaseElement (name = "element1" )
109- line1 = Line (line = [element1 ])
109+ line1 = BeamLine (line = [element1 ])
110110 assert line1 .line == [element1 ]
111111 # Extend first line with one thick element
112112 element2 = ThickElement (name = "element2" , length = 2.0 )
113113 line1 .line .extend ([element2 ])
114114 assert line1 .line == [element1 , element2 ]
115115 # Create second line with one drift element
116116 element3 = DriftElement (name = "element3" , length = 3.0 )
117- line2 = Line (line = [element3 ])
117+ line2 = BeamLine (line = [element3 ])
118118 # Extend first line with second line
119119 line1 .line .extend (line2 .line )
120120 assert line1 .line == [element1 , element2 , element3 ]
@@ -126,8 +126,8 @@ def test_yaml():
126126 # Create one thick element
127127 element2 = ThickElement (name = "element2" , length = 2.0 )
128128 # Create line with both elements
129- line = Line (line = [element1 , element2 ])
130- # Serialize the Line object to YAML
129+ line = BeamLine (line = [element1 , element2 ])
130+ # Serialize the BeamLine object to YAML
131131 yaml_data = yaml .dump (line .model_dump (), default_flow_style = False )
132132 print (f"\n { yaml_data } " )
133133 # Write the YAML data to a test file
@@ -137,11 +137,11 @@ def test_yaml():
137137 # Read the YAML data from the test file
138138 with open (test_file , "r" ) as file :
139139 yaml_data = yaml .safe_load (file )
140- # Parse the YAML data back into a Line object
141- loaded_line = Line (** yaml_data )
140+ # Parse the YAML data back into a BeamLine object
141+ loaded_line = BeamLine (** yaml_data )
142142 # Remove the test file
143143 os .remove (test_file )
144- # Validate loaded Line object
144+ # Validate loaded BeamLine object
145145 assert line == loaded_line
146146
147147
@@ -151,8 +151,8 @@ def test_json():
151151 # Create one thick element
152152 element2 = ThickElement (name = "element2" , length = 2.0 )
153153 # Create line with both elements
154- line = Line (line = [element1 , element2 ])
155- # Serialize the Line object to JSON
154+ line = BeamLine (line = [element1 , element2 ])
155+ # Serialize the BeamLine object to JSON
156156 json_data = json .dumps (line .model_dump (), sort_keys = True , indent = 2 )
157157 print (f"\n { json_data } " )
158158 # Write the JSON data to a test file
@@ -162,9 +162,9 @@ def test_json():
162162 # Read the JSON data from the test file
163163 with open (test_file , "r" ) as file :
164164 json_data = json .loads (file .read ())
165- # Parse the JSON data back into a Line object
166- loaded_line = Line (** json_data )
165+ # Parse the JSON data back into a BeamLine object
166+ loaded_line = BeamLine (** json_data )
167167 # Remove the test file
168168 os .remove (test_file )
169- # Validate loaded Line object
169+ # Validate loaded BeamLine object
170170 assert line == loaded_line
0 commit comments