Move trace parameter check to correct place

The check should be performed on all traces that are added, and it should always be done
Keysight · Nov 4, 2022 · f6f941d · f6f941d
1 parent 1cbf52e
commit f6f941d
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 5 deletions.
diff --git a/tests/test_creation.py b/tests/test_creation.py
@@ -196,6 +196,75 @@ def test_write_closed(self):
         with self.assertRaises(ValueError):
             print(trs_traces)
 
+    def test_write_different_trace_sizes(self):
+        trace_count = 100
+        sample_count = 1000
+
+        with trsfile.open(self.tmp_path, 'w', padding_mode=TracePadding.AUTO) as trs_traces:
+            trs_traces.extend([
+                Trace(
+                    SampleCoding.FLOAT,
+                    [0] * sample_count,
+                    TraceParameterMap({'LEGACY_DATA': ByteArrayParameter(i.to_bytes(8, byteorder='big'))})
+                )
+                for i in range(0, trace_count)]
+            )
+            with self.assertRaises(TypeError):
+                # The length is incorrect
+                # Should raise a Type error: The parameters of trace #0 do not match the trace set's definitions.
+                trs_traces.extend([
+                    Trace(
+                        SampleCoding.FLOAT,
+                        [0] * sample_count,
+                        TraceParameterMap({'LEGACY_DATA': ByteArrayParameter(bytes.fromhex('cafebabedeadbeef0102030405060708'))})
+                    )]
+                )
+            with self.assertRaises(TypeError):
+                # The name is incorrect
+                # Should raise a Type error: The parameters of trace #1 do not match the trace set's definitions.
+                trs_traces.extend([
+                    Trace(
+                        SampleCoding.FLOAT,
+                        [0] * sample_count,
+                        TraceParameterMap({'LEGACY_DATA': ByteArrayParameter(bytes.fromhex('0102030405060708'))})
+                    ),
+                    Trace(
+                        SampleCoding.FLOAT,
+                        [0] * sample_count,
+                        TraceParameterMap({'NEW_DATA': ByteArrayParameter(bytes.fromhex('0102030405060708'))})
+                    )]
+                )
+            with self.assertRaises(TypeError):
+                # The type is incorrect
+                # Should raise a Type error: The parameters of trace #0 do not match the trace set's definitions.
+                trs_traces.extend([
+                    Trace(
+                        SampleCoding.FLOAT,
+                        [0] * sample_count,
+                        TraceParameterMap({'LEGACY_DATA': IntegerArrayParameter([42, 74])})
+                    )]
+                )
+
+        with trsfile.open(self.tmp_path, 'w', padding_mode=TracePadding.AUTO) as trs_traces:
+            trs_traces.extend([
+                Trace(
+                    SampleCoding.FLOAT,
+                    [0] * sample_count,
+                    TraceParameterMap()
+                )
+                for i in range(0, trace_count)]
+            )
+            with self.assertRaises(TypeError):
+                # The length, data and name are incorrect
+                # Should raise a Type error: The parameters of trace #0 do not match the trace set's definitions.
+                trs_traces.extend([
+                    Trace(
+                        SampleCoding.FLOAT,
+                        [0] * sample_count,
+                        TraceParameterMap({'LEGACY_DATA': ByteArrayParameter(bytes.fromhex('cafebabedeadbeef0102030405060708'))})
+                    )]
+                )
+
     def test_read(self):
         trace_count = 100
         sample_count = 1000

diff --git a/trsfile/engine/trs.py b/trsfile/engine/trs.py
@@ -177,14 +177,18 @@ def update_headers_with_traces_metadata(self, traces: List[Trace]) -> None:
 
             # Add a TraceParameterDefinitionMap if none is present, and verify its validity if one is present
             if Header.TRACE_PARAMETER_DEFINITIONS not in self.headers:
-                if data_length > 0:
-                    headers_updates[Header.TRACE_PARAMETER_DEFINITIONS] = \
-                        TraceParameterDefinitionMap.from_trace_parameter_map(traces[0].parameters)
-            elif not traces[0].parameters.matches(self.headers[Header.TRACE_PARAMETER_DEFINITIONS]):
-                raise TypeError("The traces' parameters do not match the trace set's definitions")
+                headers_updates[Header.TRACE_PARAMETER_DEFINITIONS] = \
+                    TraceParameterDefinitionMap.from_trace_parameter_map(traces[0].parameters)
 
             headers_updates[Header.LENGTH_DATA] = data_length
 
+        # Verify that each trace confirms to the traceset's TraceParameterDefinitionMap
+        for index, trace in enumerate(traces):
+            if not trace.parameters.matches(self.headers[Header.TRACE_PARAMETER_DEFINITIONS]):
+                raise TypeError(f"The parameters of trace #{index} do not match the trace set's definitions.\n"
+                                f"Please make sure the trace parameters match those of the other traces in type, "
+                                f"size and name.")
+
         if self.headers[Header.SAMPLE_CODING] is None:
             if len(set([trace.sample_coding for trace in traces])) > 1:
                 raise TypeError('Traces have different sample coding, this is not supported in TRS files')