address review comments

Author: sid-acryl

metadata-ingestion/tests/unit/test_transform_dataset.py

@@ -654,22 +654,29 @@ def _test_owner(
             PipelineContext(run_id="test"),
         )
 
-        list(
+        record_envelops: List[RecordEnvelope] = list(
             transformer.transform(
                 [
                     RecordEnvelope(dataset, metadata={}),
+                    RecordEnvelope(record=EndOfStream(), metadata={}),
                 ]
             )
         )
 
-        mcp: MetadataChangeProposalWrapper = cast(
-            MetadataChangeProposalWrapper, transformer.handle_end_of_stream()[0]
-        )
+        assert len(record_envelops) == 3
+
+        mcp: MetadataChangeProposalWrapper = record_envelops[1].record
+
         owners_aspect = cast(OwnershipClass, mcp.aspect)
+
         owners = owners_aspect.owners
+
         owner = owners[0]
-        if expected_owner_type is not None:
-            assert owner.type == expected_owner_type
+
+        assert expected_owner_type is not None
+
+        assert owner.type == expected_owner_type
+
         assert owner.owner == expected_owner
 
     _test_owner(
@@ -678,13 +685,15 @@ def _test_owner(
             "tag_prefix": "owner:",
         },
         expected_owner="urn:li:corpuser:foo",
+        expected_owner_type=OwnershipTypeClass.TECHNICAL_OWNER,
     )
     _test_owner(
         tag="abcdef-owner:foo",
         config={
             "tag_prefix": ".*owner:",
         },
         expected_owner="urn:li:corpuser:foo",
+        expected_owner_type=OwnershipTypeClass.TECHNICAL_OWNER,
     )
     _test_owner(
         tag="owner:foo",
@@ -693,6 +702,7 @@ def _test_owner(
             "is_user": False,
         },
         expected_owner="urn:li:corpGroup:foo",
+        expected_owner_type=OwnershipTypeClass.TECHNICAL_OWNER,
     )
     _test_owner(
         tag="owner:foo",
@@ -701,6 +711,7 @@ def _test_owner(
             "email_domain": "example.com",
         },
         expected_owner="urn:li:corpuser:[email protected]",
+        expected_owner_type=OwnershipTypeClass.TECHNICAL_OWNER,
     )
     _test_owner(
         tag="owner:foo",