diff --git a/metadata-ingestion/tests/unit/test_kafka_source.py b/metadata-ingestion/tests/unit/test_kafka_source.py
index cab0a2bce7ba8c..73d2df72f72758 100644
--- a/metadata-ingestion/tests/unit/test_kafka_source.py
+++ b/metadata-ingestion/tests/unit/test_kafka_source.py
@@ -330,6 +330,7 @@ def mock_get_latest_version(subject_name: str) -> Optional[RegisteredSchema]:
             "topic2-key": "test.acryl.Topic2Key",
             "topic2-value": "test.acryl.Topic2Value",
         },
+        "ingest_schemas_as_entities": True
     }
     ctx = PipelineContext(run_id="test")
     kafka_source = KafkaSource.create(source_config, ctx)
@@ -478,8 +479,7 @@ def mock_get_latest_version(subject_name: str) -> Optional[RegisteredSchema]:
     kafka_source = KafkaSource.create(source_config, ctx)
 
     workunits = list(kafka_source.get_workunits())
-
-    assert len(workunits) == 6
+    assert len(workunits) == 2
     if ignore_warnings_on_schema_type:
         assert not kafka_source.report.warnings
     else:
@@ -622,6 +622,7 @@ def mock_get_latest_version(subject_name: str) -> Optional[RegisteredSchema]:
     kafka_source = KafkaSource.create(
         {
             "connection": {"bootstrap": "localhost:9092"},
+            "ingest_schemas_as_entities": True,
             "meta_mapping": {
                 "owner": {
                     "match": "^@(.*)",