{
  "source_file": "0930",
  "source_platform": "个人博客",
  "author": "Eugene Yan",
  "title_listed": "Writing Robust Tests for Data & Machine Learning Pipelines",
  "keywords": "测试,数据管道,ML",
  "domain": "数据处理",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://eugeneyan.com/writing/testing-pipelines/",
  "publish_date": "2022-10",
  "corpus_stage": "推理部署",
  "notes": "数据与ML管道健壮测试",
  "id": "b04dabee6e8d",
  "collected_at": "2026-10-02T14:17:03",
  "title": "Writing Robust Tests for Data & Machine Learning Pipelines",
  "body_markdown": "eugeneyan\nStart Here\nWriting\nSpeaking\nPrototyping\nAbout\nWriting Robust Tests for Data & Machine Learning Pipelines\n[\nengineering\nmachinelearning\nproduction\n🩷\n]\n· 19 min read\nRecently, I’ve been pondering on how data and machine learning pipelines are tested, especially why some tests break more often than others. And it’s not because the new code is wrong; often, the new code is correct but the tests break anyway and need to be updated.\nHere, I dig into why certain tests break—incorrectly—more than others and try to find a better way to test pipelines. We’ll start with a simple pipeline and test it via unit, schema, and integration tests. Then, I’ll introduce new data and logic, observe how tests break, and draw patterns from it. Finally, I’ll suggest how to make pipeline testing less brittle.\nOverview of testing scopes: unit, integration, functional, etc.\nAn example pipeline: behavioral logs -> batch inference output\nWriting tests for our pipeline: unit, schema, integration\nAdding new data (visible impressions) and logic to our pipeline\nThe additive vs. retroactive impact of new data and logic on tests\nRelated concepts: Test validity, granularity, and scope\nSuggestions for more robust pipeline tests\nSmaller testing scopes = Shorter feedback loops\nBefore we jump into it, here’s a brief overview of various testing scopes (an Atlassian\narticle\nlists as many as seven of them). If we’re building a recommendation system pipeline, here’s how the various scopes map to concrete tests:\nUnit: Test individual methods and classes\nIntegration: Test pipeline integration points, such as between data processing and model training pipelines, or checking that trained models are deployed correctly\nFunctional: Test that business requirements are met (e.g., excluding erotica, rolling up movies from the same series, updating recs based on in-session clicks, etc.)\nEnd-to-end: Test that the pipeline output/endpoint integrates correctly with the site and recommendations are served correctly based on user ID or behavior\nAcceptance: Visual inspection and quality assurance of live recommendation slate\nPerformance: Load testing and assessing latency and throughput\nSmoke: Handcrafted URLs to trigger recommendations on commonly QA-ed items\nA/B: Evaluate the impact of pipeline or model changes\nTests increase in scope as we go down the list. Bigger scopes lead to longer feedback loops: A unit test runs in milliseconds while an A/B test requires time to implement the change plus whatever period the A/B test needs to run for.\nHere, we’ll focus on tests with the shortest feedback loops: unit tests and integration tests.\nSpecifically, we’ll discuss tests that can be run offline, either on our local machine or a lightweight cloud instance. (Although integration tests can also run in a CI environment that calls other services or queries a database, we’ll exclude those.)\nAlso, though data and machine learning pipelines often have additional tests on data quality (e.g.,\nGreat Expectations\n,\nDeequ\n), model evaluation (e.g.,\nscikit-learn metrics\n), and\nmodel behavioral checks\n(e.g.,\nCheckList\n,\nRecList\n), these tend to run in production with each new batch of data or model refresh. Thus, we’ll not discuss them here.\nExample pipeline: Behavioral logs -> inference output\nAt a high level, our example pipeline (i) consumes request-level behavioral logs, (ii) transforms them into item-level events, (iii) trains a naive model that estimates click-through rate (CTR), and (iv) performs batch inference on a list of items.\nHere’s what the\nbehavioral logs\nlook like; this is also the sample data we’ll use in our tests. They track each time a recommendation slate is requested and the items that have been impressed or clicked. (The raw form of such logs is usually unstructured but in this pipeline they’ve been parsed and stored in a nice tabular format.)\n+--------------+--------------------------+--------------+--------------+\n|\nrequest_id\n|\nimpressions\n|\nevent_item\n|\nevent_type\n|\n|--------------+--------------------------+--------------+--------------|\n|\nr1\n|\n[\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n]\n|\n|\nimpress\n|\n|\nr2\n|\n[\n'i2'\n,\n'i3'\n,\n'i4'\n,\n'i1'\n]\n|\ni2\n|\nclick\n|\n|\nr3\n|\n[\n'i3'\n,\n'i4'\n,\n'i1'\n,\n'i2'\n]\n|\n|\nimpress\n|\n|\nr4\n|\n[\n'i4'\n,\n'i1'\n,\n'i2'\n,\n'i3'\n]\n|\ni3\n|\nclick\n|\n|\nr5\n|\n[\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n]\n|\n|\nimpress\n|\n+--------------+--------------------------+--------------+--------------+\nA brief description of each column:\nrequest_id\n: ID of recommendation slate request\nimpressions\n: IDs of impressed items\nevent_item\n: ID of clicked items; empty otherwise\nevent_type\n: Type of event (i.e., impress or click)\nFor request\nr1\n, four items (\ni1\n,\ni2\n,\ni3\n,\ni4\n) were impressed but none were clicked. In\nr2\n, four items (\ni2\n,\ni3\n,\ni4\n,\ni1\n) were impressed and\ni2\n—the item in the first position—was clicked.\nTo make it easier to estimate item CTR, we want to transform the\nrequest-level\nbehavioral logs into\nitem-level\nevents where each impression and click has its own row. In addition, we also want to get the position of the item that was impressed or clicked.\nHere’s what the\nevents\ntable looks like for\nr2\nafter the transformation:\n+--------------+--------+------------+--------------+\n|\nrequest_id\n|\nitem\n|\nposition\n|\nevent_type\n|\n|--------------+--------+------------+--------------|\n|\nr2\n|\ni2\n|\n1\n|\nimpress\n|\n|\nr2\n|\ni3\n|\n2\n|\nimpress\n|\n|\nr2\n|\ni4\n|\n3\n|\nimpress\n|\n|\nr2\n|\ni1\n|\n4\n|\nimpress\n|\n|\nr2\n|\ni2\n|\n1\n|\nclick\n|\n+--------------+--------+------------+--------------+\nNext, we train a naive model that learns historical CTR and estimates item-level CTR. This model is then used to perform batch inference on a list of item IDs.\nHere’s what the\nbatch inference output\nlooks like. Notice that the expected CTR for\ni5\nis -1. This is because we only have behavioral logs for\ni1\nto\ni4\n—thus,\ni5\nis a new or cold-start item. Since our model can’t estimate CTR for\ni5\n, it returns a predefined null value (-1).\n+-----------+----------------+\n|\nitem_id\n|\nexpected_ctr\n|\n|-----------+----------------|\n|\ni1\n|\n0\n|\n|\ni2\n|\n0.2\n|\n|\ni3\n|\n0.2\n|\n|\ni4\n|\n0\n|\n|\ni5\n|\n-\n1\n|\n+-----------+----------------+\nSo far, I’ve focused on sharing sample data at various stages in our example pipeline. If you prefer to look at code to help you understand better, please refer to\npipeline.py\nhere\n. Note that the code is deliberately kept simple as its implementation isn’t key in this write-up.\nImplementation tests: Unit, Schema, Integration\nTo test our pipeline, we’ll add (i) row-level unit tests, (ii) column-level unit tests, (iii) table-level unit tests, (iv) schema tests, and (v) integration tests. We use basic Pandas (\nassert_frame_equal\n) and PyTest (\nfixtures\n) functionality. See test fixtures on\nGitHub\n.\nLet’s start with the smallest form,\nrow-level unit tests\n, where we test a method that takes an input (row) and returns a single value. We’ll test\nget_click_position\nwhich takes an array of impressions and the clicked item, and returns the position of the clicked item. Writing the test for this is straightforward and we can parametrize our test with various inputs and expected outputs.\ndef\nget_click_position\n(\nimpressions\n:\nList\n[\nstr\n],\nclick\n:\nstr\n)\n->\nint\n:\n\"\"\"Returns the position of the clicked item based on the array of impressions.\"\"\"\ntry\n:\nreturn\nimpressions\n.\nindex\n(\nclick\n)\n+\n1\nexcept\nValueError\n:\nreturn\n-\n1\n# Unit test: Row level\n@\npytest\n.\nmark\n.\nparametrize\n(\n'impressions,click,expected'\n,\n[([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\n'i1'\n,\n1\n),\n([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\n'i3'\n,\n3\n),\n([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\nNone\n,\n-\n1\n),\n([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\n'NA'\n,\n-\n1\n)])\ndef\ntest_get_click_position\n(\nimpressions\n,\nclick\n,\nexpected\n):\nassert\nget_click_position\n(\nimpressions\n,\nclick\n)\n==\nexpected\nNext, we add\ncolumn-level unit tests\nfor methods that take a column or table and return a column. The\nget_impress_positions\nmethod takes a table of request IDs and impression events and returns a column of each impression event’s position. Because the expected output is a column instead of a single value, the test starts to get unwieldy and it can be tedious to define expected outputs.\ndef\nget_impress_positions\n(\ndf\n:\npd\n.\nDataFrame\n)\n->\npd\n.\nSeries\n:\n\"\"\"Returns a column of impression positions.\n\n       Note: This method assumes that impressions are sorted in ascending order of their position.\n    \"\"\"\npositions\n=\ndf\n.\ngroupby\n(\n'request_id'\n).\ncumcount\n()\n+\n1\npositions\n=\npositions\n.\nreset_index\n(\ndrop\n=\nTrue\n)\nreturn\npositions\n# Unit test: Column level\ndef\ntest_get_impress_position\n(\nlogs\n):\nimpress_events\n=\nlogs\n.\nexplode\n(\n'impressions'\n)\nimpress_positions\n=\nget_impress_positions\n(\nimpress_events\n)\n# Since each impress log has 4 items, expect impression positions to be five sets of 1-4\npd\n.\ntesting\n.\nassert_series_equal\n(\nimpress_positions\n,\npd\n.\nSeries\n([\n1\n,\n2\n,\n3\n,\n4\n]\n*\n5\n))\nNext, we add\ntable-level unit tests\nfor methods that take a table, do some aggregation or filtering, and return another table. The\naggregate_events\nmethod takes a table of impression and click events and aggregates them at the item level like below.\n+--------+---------+-----------+\n|\nitem\n|\nclick\n|\nimpress\n|\n|--------+---------+-----------|\n|\ni1\n|\n0\n|\n5\n|\n|\ni2\n|\n1\n|\n5\n|\n|\ni3\n|\n1\n|\n5\n|\n|\ni4\n|\n0\n|\n5\n|\n+--------+---------+-----------+\nHere’s the method and its unit test. Table-level unit tests are even more clunky as we need to define entire tables as expected output.\ndef\naggregate_events\n(\nevents\n:\npd\n.\nDataFrame\n)\n->\npd\n.\nDataFrame\n:\n\"\"\"Returns a table of items and their aggregated impressions and clicks.\"\"\"\nevents_agg\n=\nevents\n.\npivot_table\n(\nindex\n=\n[\n'item'\n],\ncolumns\n=\n[\n'event_type'\n],\nvalues\n=\n[\n'request_id'\n],\naggfunc\n=\n'count'\n,\nfill_value\n=\n0\n)\nevents_agg\n.\ncolumns\n=\nevents_agg\n.\ncolumns\n.\ndroplevel\n()\nevents_agg\n=\nevents_agg\n.\nreset_index\n()\n# Clear index name\nevents_agg\n=\nevents_agg\n.\nrename_axis\n(\nNone\n,\naxis\n=\n1\n)\nreturn\nevents_agg\n# Unit test: Table level\ndef\ntest_aggregate_events\n(\nevents\n):\nresult\n=\naggregate_events\n(\nevents\n)\narr\n=\n[[\n'i1'\n,\n0\n,\n5\n],\n[\n'i2'\n,\n1\n,\n5\n],\n[\n'i3'\n,\n1\n,\n5\n],\n[\n'i4'\n,\n0\n,\n5\n]]\ncols\n=\n[\n'item'\n,\n'click'\n,\n'impress'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\nWe also add\nschema tests\nat key pipeline junctures. Here, we check for expected column names and datatypes after transforming the request-level behavioral logs to item-level events. This test checks for a\nminimum\nset of columns: the pipeline can be updated without failing this test so long as the minimum set of columns is present with the right datatypes.\n# Schema check: Columns and datatypes\ndef\ntest_events_schema\n(\nevents\n):\nexpected_cols\n=\n{\n'request_id'\n:\nnp\n.\nobject\n,\n'item'\n:\nnp\n.\nobject\n,\n'position'\n:\nnp\n.\nint64\n,\n'event_type'\n:\nnp\n.\nobject\n}\n# Check all expected columns are present\nassert\nlen\n(\nset\n(\nexpected_cols\n.\nkeys\n())\n-\nset\n(\nevents\n.\ncolumns\n))\n==\n0\n,\n\\\nf\n'\n{\nset\n(\nexpected_cols\n.\nkeys\n())\n-\nset\n(\nevents\n.\ncolumns\n)\n}\ncolumns missing!'\n# Check all column data types are correct\nfor\ncol\n,\ndtype\nin\nexpected_cols\n.\nitems\n():\nassert\nevents\n.\ndtypes\n[\ncol\n]\n==\ndtype\n,\n\\\nf\n'Expected column\n{\ncol\n}\nto be of\n{\ndtype\n}\ntype but found\n{\nevents\n.\ndtypes\n[\ncol\n]\n}\ntype!'\nFinally, we add\nintegration tests\nthat take a fixed sample input, run it through the pipeline, and compare the result with the expected output. Below, we test that behavioral logs are transformed into events correctly.\n# Integration test: Input logs to aggregated events\ndef\ntest_feature_pipeline\n(\nlogs\n):\nimpress_logs\n=\nlogs\nclick_logs\n=\nlogs\n[\nlogs\n[\n'event_type'\n]\n==\n'click'\n]\nimpress_events\n=\nget_impress_item_and_pos\n(\nimpress_logs\n)\nclick_events\n=\nget_click_item_and_pos\n(\nclick_logs\n)\nsample_events\n=\npd\n.\nconcat\n([\nimpress_events\n,\nclick_events\n])\nresult\n=\naggregate_events\n(\nsample_events\n)\narr\n=\n[[\n'i1'\n,\n0\n,\n5\n],\n[\n'i2'\n,\n1\n,\n5\n],\n[\n'i3'\n,\n1\n,\n5\n],\n[\n'i4'\n,\n0\n,\n5\n]]\ncols\n=\n[\n'item'\n,\n'click'\n,\n'impress'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\nWe also test that our model learns and returns the expected CTR estimates correctly.\n# Integration test: Aggregated events to batch inference\ndef\ntest_model_pipeline\n(\nevents\n,\nmodel\n,\nitems\n):\nmodel\n=\nmodel\n.\nfit\n(\nevents\n)\nresult\n=\nmodel\n.\nbatch_predict\n(\nitems\n)\narr\n=\n[[\n'i1'\n,\n0.0\n],\n[\n'i2'\n,\n0.2\n],\n[\n'i3'\n,\n0.2\n],\n[\n'i4'\n,\n0.0\n],\n[\n'i5'\n,\n-\n1.0\n]]\ncols\n=\n[\n'item_id'\n,\n'expected_ctr'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\nIntegration tests are likely the most onerous of the lot as they rely on multiple data transforms and training steps working correctly. Thus, they can become unmanageable if the pipeline is updated frequently and they break just as often. Nonetheless, I find them helpful in ensuring the pipeline works as intended. Their expected output also helps explain how the pipeline should work.\nHere’s the result of running our\ntests\n.\n>\npytest\n-v\ntests/test_pipeline.py\n--cov\n=\nsrc\n--cov-report\n=\nterm-missing\n\nName              Stmts   Miss  Cover   Missing\n-----------------------------------------------\nsrc/__init__.py       0      0   100%\nsrc/pipeline.py      42      0   100%\n-----------------------------------------------\nTOTAL                42      0   100%\nAside: In this simple example, the pipeline code fits in a single script. In an industrial setting, the pipeline may span multiple repos and infra such as Spark (Scala, SQL) and SageMaker (Python, Docker). It’s possible to set up and run offline tests across multiple repos, but depending on your application and scale, it may not be worth the effort.\nAdding new data or logic = Updating pipeline code\nIn the current pipeline, the logged impressions are\nserver-side\nimpressions: All items in the recommendation slate are logged as impressed on customers\neven if the customer has not actually seen the item\n.\nFor example, on Netflix, only the first six items in a recommendation slate are visible on the screen; I need to side scroll to view the rest of the items. With server-side impressions, even if I didn’t side scroll, all items in the recommendation slate would be logged as impressed. This leads to a discrepancy between what we think the customer saw vs. what they really saw, especially on smaller screens such as mobile phones.\nOnly a handful of items in a recommendation slate are visible if users don't side scroll.\nTo address this, we add additional data to our pipeline:\nclient-side\nimpressions. This data tracks which items were actually visible on the screen. Thus, if I didn’t side scroll on the recommendation slate, only the first six items would be logged as client-side impressions.\nHere’s what the\nupdated behavioral logs\nlook like. We have a new column for visible impressions. (But client-side data’s not perfect which explains why\nr3\nhas a null value.)\n+--------------+--------------------------+--------------------------+--------------+--------------+\n|\nrequest_id\n|\nimpressions\n|\nimpressions_visible\n|\nevent_item\n|\nevent_type\n|\n|--------------+--------------------------+--------------------------+--------------+--------------|\n|\nr1\n|\n[\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n]\n|\n[\n'i1'\n,\n'i2'\n,\n'i3'\n]\n|\n|\nimpress\n|\n|\nr2\n|\n[\n'i2'\n,\n'i3'\n,\n'i4'\n,\n'i1'\n]\n|\n[\n'i2'\n,\n'i3'\n]\n|\ni2\n|\nclick\n|\n|\nr3\n|\n[\n'i3'\n,\n'i4'\n,\n'i1'\n,\n'i2'\n]\n|\n|\n|\nimpress\n|\n|\nr4\n|\n[\n'i4'\n,\n'i1'\n,\n'i2'\n,\n'i3'\n]\n|\n[\n'i4'\n,\n'i1'\n,\n'i2'\n,\n'i3'\n]\n|\ni3\n|\nclick\n|\n|\nr5\n|\n[\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n]\n|\n[\n'i1'\n,\n'i2'\n]\n|\n|\nimpress\n|\n+--------------+--------------------------+--------------------------+--------------+--------------+\nWe’ll want to update our pipeline to only consider impressions that were visible via some basic logic. Thus, we add methods to update impressions (at the row level) before updating the impressions column (at the table level).\ndef\nget_updated_impressions\n(\nimpressions\n:\nList\n[\nstr\n],\nimpressions_visible\n:\nList\n[\nstr\n])\n->\nList\n[\nstr\n]:\n\"\"\"Returns an array of updated impressions based on logs of visible impressions.\"\"\"\nif\nisinstance\n(\nimpressions_visible\n,\nlist\n)\nand\nlen\n(\nimpressions\n)\n>=\nlen\n(\nimpressions_visible\n):\nreturn\nimpressions_visible\nreturn\nimpressions\ndef\nupdate_impression_col\n(\ndf\n:\npd\n.\nDataFrame\n)\n->\npd\n.\nDataFrame\n:\n\"\"\"Returns a table of logs where impressions is updated to consider visible impressions.\"\"\"\n_df\n=\ndf\n.\ncopy\n()\n_df\n[\n'impressions'\n]\n=\n_df\n.\napply\n(\nlambda\nx\n:\nget_updated_impressions\n(\nx\n[\n'impressions'\n],\nx\n[\n'impressions_visible'\n]),\naxis\n=\n1\n)\n_df\n=\n_df\n.\ndrop\n(\ncolumns\n=\n[\n'impressions_visible'\n])\nreturn\n_df\nWhen we try to run our existing tests again with our new\ninput\nand\npipeline code\n, not counting the parameterization,\n4/6 tests fail because of the new data and logic\n.\n> pytest -v tests/test_pipeline.py --cov=src --cov-report=term-missing\n\ntests/test_pipeline.py::test_get_click_position[impressions0-i1-1] PASSED       [ 11%]\ntests/test_pipeline.py::test_get_click_position[impressions1-i3-3] PASSED       [ 22%]\ntests/test_pipeline.py::test_get_click_position[impressions2-None--1] PASSED    [ 33%]\ntests/test_pipeline.py::test_get_click_position[impressions3-NA--1] PASSED      [ 44%]\ntests/test_pipeline.py::test_get_impress_position FAILED                        [ 55%]\ntests/test_pipeline.py::test_aggregate_events FAILED                            [ 66%]\ntests/test_pipeline.py::test_events_schema PASSED                               [ 77%]\ntests/test_pipeline.py::test_feature_pipeline FAILED                            [ 88%]\ntests/test_pipeline.py::test_model_pipeline FAILED                              [100%]\nAdding new data/logic: Additive vs. retroactive impacts\nSince we’ve updated our pipeline, we’ll need to\nadd tests\nfor the\nget_updated_impressions\nmethod which is easy to parameterize.\ndef\nget_updated_impressions\n(\nimpressions\n:\nList\n[\nstr\n],\nimpressions_visible\n:\nList\n[\nstr\n])\n->\nList\n[\nstr\n]:\n\"\"\"Returns an array of updated impressions based on logs of visible impressions.\"\"\"\nif\nisinstance\n(\nimpressions_visible\n,\nlist\n)\nand\nlen\n(\nimpressions\n)\n>=\nlen\n(\nimpressions_visible\n):\nreturn\nimpressions_visible\nreturn\nimpressions\n# Unit test: Row level (added)\n@\npytest\n.\nmark\n.\nparametrize\n(\n'impressions,impressions_visible,updated_impressions'\n,\n[([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\n[\n'i1'\n,\n'i2'\n,\n'i3'\n],\n[\n'i1'\n,\n'i2'\n,\n'i3'\n]),\n([\n'i1'\n,\n'i2'\n],\n[\n'i1'\n,\n'i2'\n,\n'i3'\n],\n[\n'i1'\n,\n'i2'\n]),\n([\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n],\nNone\n,\n[\n'i1'\n,\n'i2'\n,\n'i3'\n,\n'i4'\n])])\ndef\ntest_get_updated_impressions\n(\nimpressions\n,\nimpressions_visible\n,\nupdated_impressions\n):\nassert\nget_updated_impressions\n(\nimpressions\n,\nimpressions_visible\n)\n==\nupdated_impressions\nWhat about the existing tests?\nExisting row-level unit tests are unchanged.\nBecause the tested methods only depend on row-level input, the test is invariant to new data. Barring incorrect data, the logic to compute click position isn’t coupled to whether the impression is visible or not.\nExisting column-level unit tests need to be updated\nbecause we’re now excluding non-visible impressions. The new expected output needs to be manually crafted.\n# BEFORE - Unit test: Column level\ndef\ntest_get_impress_position\n(\nlogs\n):\nimpress_events\n=\nlogs\n.\nexplode\n(\n'impressions'\n)\nimpress_positions\n=\nget_impress_positions\n(\nimpress_events\n)\n# Since each impress log has 4 items, expect impression positions to be five sets of 1-4\npd\n.\ntesting\n.\nassert_series_equal\n(\nimpress_positions\n,\npd\n.\nSeries\n([\n1\n,\n2\n,\n3\n,\n4\n]\n*\n5\n))\n# AFTER - Unit test: Column level (updated)\ndef\ntest_get_impress_position\n(\nlogs\n):\nlogs\n=\nupdate_impression_col\n(\nlogs\n)\nimpress_events\n=\nlogs\n.\nexplode\n(\n'impressions'\n)\nimpress_positions\n=\nget_impress_positions\n(\nimpress_events\n)\n# Update impress positions (too brittle?)\npd\n.\ntesting\n.\nassert_series_equal\n(\nimpress_positions\n,\npd\n.\nSeries\n([\n1\n,\n2\n,\n3\n,\n1\n,\n2\n,\n1\n,\n2\n,\n3\n,\n4\n,\n1\n,\n2\n,\n3\n,\n4\n,\n1\n,\n2\n]))\nExisting table-level unit tests need to be updated\nbecause the filtered impressions affect aggregation results.\n# BEFORE - Unit test: Table level\ndef\ntest_aggregate_events\n(\nevents\n):\nresult\n=\naggregate_events\n(\nevents\n)\narr\n=\n[[\n'i1'\n,\n0\n,\n5\n],\n[\n'i2'\n,\n1\n,\n5\n],\n[\n'i3'\n,\n1\n,\n5\n],\n[\n'i4'\n,\n0\n,\n5\n]]\ncols\n=\n[\n'item'\n,\n'click'\n,\n'impress'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\n# AFTER - Unit test: Table level (updated)\ndef\ntest_aggregate_events\n(\nevents\n):\nresult\n=\naggregate_events\n(\nevents\n)\narr\n=\n[[\n'i1'\n,\n0\n,\n4\n],\n# Updated impress count\n[\n'i2'\n,\n1\n,\n5\n],\n[\n'i3'\n,\n1\n,\n4\n],\n# Updated impress count\n[\n'i4'\n,\n0\n,\n2\n]]\n# Updated impress count\ncols\n=\n[\n'item'\n,\n'click'\n,\n'impress'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\nExisting schema tests are unchanged.\nRegardless of how we transform behavioral logs, the schema for the minimum set of columns should stay constant.\nExisting integration tests need to be updated\nfor the same reason why column/table-level unit tests need to be updated. Any updates to data or logic will affect intermediate data throughout the pipeline as well as the batch-inference output at the end of the pipeline.\n# BEFORE - Integration test: Aggregated events to batch inference\ndef\ntest_model_pipeline\n(\nevents\n,\nmodel\n,\nitems\n):\nmodel\n=\nmodel\n.\nfit\n(\nevents\n)\nresult\n=\nmodel\n.\nbatch_predict\n(\nitems\n)\narr\n=\n[[\n'i1'\n,\n0.0\n],\n[\n'i2'\n,\n0.2\n],\n[\n'i3'\n,\n0.2\n],\n[\n'i4'\n,\n0.0\n],\n[\n'i5'\n,\n-\n1.0\n]]\ncols\n=\n[\n'item_id'\n,\n'expected_ctr'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\n# AFTER - Integration test: Aggregated events to batch inference (updated)\ndef\ntest_model_pipeline\n(\nevents\n,\nmodel\n,\nitems\n):\nmodel\n=\nmodel\n.\nfit\n(\nevents\n)\nresult\n=\nmodel\n.\nbatch_predict\n(\nitems\n)\narr\n=\n[[\n'i1'\n,\n0.0\n],\n[\n'i2'\n,\n0.2\n],\n[\n'i3'\n,\n0.25\n],\n# Update expected CTR\n[\n'i4'\n,\n0.0\n],\n[\n'i5'\n,\n-\n1.0\n]]\ncols\n=\n[\n'item_id'\n,\n'expected_ctr'\n]\nexpected\n=\npd\n.\nDataFrame\n(\narr\n,\ncolumns\n=\ncols\n)\npd\n.\ntesting\n.\nassert_frame_equal\n(\nresult\n,\nexpected\n)\nTo recap:\nUnit tests (row-level): Existing tests unchanged; add tests for new methods.\nUnit tests (col-level): Update existing tests.\nUnit tests (table-level): Update existing tests.\nSchema test: Existing tests unchanged; optionally add tests for new columns.\nIntegration tests: Update existing tests.\nVia this simple example, we see that\nthe effect of new data or logic is (i) additive for row-level unit tests and schema tests but (ii) retroactive for column/table-level unit tests and integration tests.\nFor the former, we simply add new tests while keeping existing tests unchanged. For the latter, we need to update most existing tests—they’re\nbrittle to changes in pipeline data and logic\n.\nAlso, considering how the latter tests are more time-consuming to write and update, this explains why we spend more time on them than on writing feature code.\nIn the example pipeline and tests above, updating the tests wasn’t too onerous as we only had to update the code generating the input and expected output. Furthermore, the tables were small with few rows and columns. However, if we have more data in a format that’s less accessible in an IDE (e.g., Parquet), it would be a pain to update the input and expected output across many rows and columns.\nOne way to address this is data generation and property-based testing.\nData generation involves generating fake data based on specifications such as numerical distributions and valid categorical values. Libraries such as\nFaker\nand\nHypothesis\nhelp with this.\nTo apply property-based testing to our pipeline, we would (i) generate a lot of data, (ii) put the data through the pipeline, and (iii) check that the output meets the predefined properties. (In contrast, what we’ve been discussing is example-based testing: Given an input, check for the fixed output.) Libraries that support property-based testing include:\nHypothesis\n,\nPandera\n(via Hypothesis and with additional schema checks), and\nGiskard\n(perturbations on text, image, and tabular input and tests on model variance).\nProperty-based testing seems like a promising approach. I used Faker for it a while back. Nonetheless, I’m uncertain about its ability to test column and table-level data transforms. Furthermore, property-based testing is\nfuzzy\n: We pass fuzzed, generated data to the code and test if it fails. While this helps to test edge cases at scale, I’m not convinced it can easily be applied to test for specific logic such as checking that items are ranked correctly before and after business logic (e.g., diversification, boosting).\nRelated: Test validity, granularity, and scope\nIf a test breaks and has to be updated frequently, is it a\nvalid\ntest? Maybe not. A test should protect against regression—if we have to change it often, it’s likely not working correctly.\nFinding the right\ngranularity\nfor a test can be challenging but it makes testing more robust. By testing for exact values in the integration tests above, I overly coupled expected output with implementation details. Instead, I probably should have tested on more coarse-grained values such as the number of rows, columns, and unique item IDs. It would have made those integration tests more robust and not break when new data was added.\nFinally, we should test as early as possible (i.e., at the smallest\nscopes\n) and in the most generic way. Row-level unit tests do this by considering only a single row’s input and output. With comprehensive row-level unit tests, we can rely less on integration tests.\nMany row-level and schema tests, a handful of the rest\nTo sum up, row-level unit tests and schema tests are less effort to write and maintain. They’re also robust to new data and business logic. Furthermore, having comprehensive row-level tests reduces reliance on higher-level tests. I tend to apply them generously.\nIn contrast, column/table-level unit tests and integration tests tend to be brittle. Finding the right granularity for a test can be challenging but makes testing more robust. Despite the downsides, these tests provide an irreplaceable safety harness though they should be used sparingly lest they become too cumbersome to maintain. I’ve come to depend on the assurance they provide but will probably use less of them in the future.\nWhat’s your take on testing data or machine learning pipelines? What approaches have you found effective? I would love to hear from you!\nFurther Reading\nThe Different Types of Software Testing\nWhat are the Differences between Unit Tests, Integration Tests, Smoke Tests, etc?\nEffective Testing for Spark Programs\nThe Mechanics of Testing Large Data Pipelines\nTest Strategies for Data Processing Pipelines\nEngineering Data Quality\nTesting Features with PyTest\nIntro to Property-based Testing in Python\nWhat is Property Based Testing?\nTesting in Microsoft Recommenders\nThanks\nto\nDavid Said\n,\nJim Dowling\n,\nShreya Shankar\n,\nMike Lam\n, and\nLars Albertsson\nfor discussions on this topic. OG image by\nSam Loyd\non\nUnsplash\n.\nIf you found this useful, please cite this write-up as:\nYan, Ziyou. (Sep 2022). Writing Robust Tests for Data & Machine Learning Pipelines. eugeneyan.com.\n        https://eugeneyan.com/writing/testing-pipelines/.\nor\n@article{yan2022testing,\n  title   = {Writing Robust Tests for Data & Machine Learning Pipelines},\n  author  = {Yan, Ziyou},\n  journal = {eugeneyan.com},\n  year    = {2022},\n  month   = {Sep},\n  url     = {https://eugeneyan.com/writing/testing-pipelines/}\n}\nShare on:\nBrowse related tags:\n[\nengineering\nmachinelearning\nproduction\n🩷\n]\nor\nSearch\n« Simplicity is An Advantage but Sadly Complexity Sells Better\nRecSys 2022 Keynote - Is the Juice Worth the Squeeze? »\nJoin\n11,800+\nreaders getting updates on machine learning, RecSys, LLMs, and engineering.\nGet email updates",
  "metadata": {
    "title": "Writing Robust Tests for Data & Machine Learning Pipelines",
    "description": "Or why I should write fewer integration tests.",
    "keywords": null,
    "author": "Eugene Yan",
    "og:site_name": "eugeneyan.com",
    "og:title": "Writing Robust Tests for Data & Machine Learning Pipelines",
    "og:type": "article",
    "og:description": "Or why I should write fewer integration tests.",
    "og:url": "https://eugeneyan.com/writing/testing-pipelines/",
    "og:image": "https://eugeneyan.com/assets/og_image/testing-pipelines.jpg",
    "twitter:card": "summary_large_image",
    "twitter:domain": "eugeneyan.com",
    "twitter:title": "Writing Robust Tests for Data & Machine Learning Pipelines",
    "twitter:description": "Or why I should write fewer integration tests.",
    "twitter:image": "https://eugeneyan.com/assets/og_image/testing-pipelines.jpg",
    "article:published_time": "2022-09-04T00:00:00+00:00",
    "article:author": "https://eugeneyan.com/about/",
    "article:section": "posts",
    "article:tag": "🩷"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/b04dabee6e8d/code/body-000.txt",
      "chars": 666,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-001.txt",
      "chars": 486,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-002.txt",
      "chars": 279,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-003.txt",
      "chars": 707,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-004.txt",
      "chars": 695,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-005.txt",
      "chars": 264,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-006.txt",
      "chars": 846,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-007.txt",
      "chars": 675,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-008.txt",
      "chars": 626,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-009.txt",
      "chars": 435,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-010.txt",
      "chars": 330,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-011.txt",
      "chars": 909,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-012.txt",
      "chars": 776,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-013.txt",
      "chars": 855,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-014.txt",
      "chars": 890,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-015.txt",
      "chars": 1012,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-016.txt",
      "chars": 786,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-017.txt",
      "chars": 923,
      "language": "text"
    },
    {
      "file": "attachments/b04dabee6e8d/code/body-018.txt",
      "chars": 257,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/icon-search.svg",
      "bytes": 425,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-feeb4bd87b8704a903b2.svg",
      "src": "https://eugeneyan.com/assets/icon-search.svg",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/netflix-recommendation-slate.webp",
      "bytes": 30214,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-0caefa46a0335b8f94e9.webp",
      "src": "https://eugeneyan.com/assets/netflix-recommendation-slate.webp",
      "alt": "Image",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/icon-twitter.svg",
      "bytes": 892,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-0083a7581b07e5f40e7b.svg",
      "src": "https://eugeneyan.com/assets/icon-twitter.svg",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/icon-linkedin.svg",
      "bytes": 559,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-f7a16c3e46e5838cec4d.svg",
      "src": "https://eugeneyan.com/assets/icon-linkedin.svg",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/bluesky.svg",
      "bytes": 745,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-be754e8ecdaf6de94f55.svg",
      "src": "https://eugeneyan.com/assets/bluesky.svg",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/icon-facebook.svg",
      "bytes": 1123,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-5d55b59733923cf14223.svg",
      "src": "https://eugeneyan.com/assets/icon-facebook.svg",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://eugeneyan.com/assets/icon-mail.svg",
      "bytes": 841,
      "reused": false,
      "local_path": "attachments/b04dabee6e8d/images/body-b179147617228dfafb9f.svg",
      "src": "https://eugeneyan.com/assets/icon-mail.svg",
      "alt": "",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://eugeneyan.com/",
      "label": "eugeneyan",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/start-here/",
      "label": "Start Here",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/writing/",
      "label": "Writing",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/speaking/",
      "label": "Speaking",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/prototyping/",
      "label": "Prototyping",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/about/",
      "label": "About",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/search/",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/tag/engineering/",
      "label": "engineering",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/tag/machinelearning/",
      "label": "machinelearning",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/tag/production/",
      "label": "production",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/tag/🩷/",
      "label": "🩷",
      "type": "a"
    },
    {
      "href": "https://www.atlassian.com/continuous-delivery/software-testing/types-of-software-testing",
      "label": "article",
      "type": "a"
    },
    {
      "href": "https://greatexpectations.io/",
      "label": "Great Expectations",
      "type": "a"
    },
    {
      "href": "https://github.com/awslabs/deequ",
      "label": "Deequ",
      "type": "a"
    },
    {
      "href": "https://scikit-learn.org/stable/modules/model_evaluation.html",
      "label": "scikit-learn metrics",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/writing/testing-ml/",
      "label": "model behavioral checks",
      "type": "a"
    },
    {
      "href": "https://github.com/marcotcr/checklist",
      "label": "CheckList",
      "type": "a"
    },
    {
      "href": "https://reclist.io/",
      "label": "RecList",
      "type": "a"
    },
    {
      "href": "https://github.com/eugeneyan/testing-pipelines/blob/main/src/pipeline.py",
      "label": "here",
      "type": "a"
    },
    {
      "href": "https://pandas.pydata.org/docs/reference/api/pandas.testing.assert_frame_equal.html",
      "label": "assert_frame_equal",
      "type": "a"
    },
    {
      "href": "https://docs.pytest.org/en/6.2.x/fixture.html",
      "label": "fixtures",
      "type": "a"
    },
    {
      "href": "https://github.com/eugeneyan/testing-pipelines/blob/main/tests/test_pipeline.py",
      "label": "GitHub",
      "type": "a"
    },
    {
      "href": "https://github.com/eugeneyan/testing-pipelines/blob/main/tests/test_pipeline_v2.py",
      "label": "input",
      "type": "a"
    },
    {
      "href": "https://github.com/eugeneyan/testing-pipelines/blob/main/src/pipeline_v2.py",
      "label": "pipeline code",
      "type": "a"
    },
    {
      "href": "https://faker.readthedocs.io/en/master/",
      "label": "Faker",
      "type": "a"
    },
    {
      "href": "https://hypothesis.readthedocs.io/en/latest/",
      "label": "Hypothesis",
      "type": "a"
    },
    {
      "href": "https://pandera.readthedocs.io/en/stable/",
      "label": "Pandera",
      "type": "a"
    },
    {
      "href": "https://github.com/Giskard-AI/giskard",
      "label": "Giskard",
      "type": "a"
    },
    {
      "href": "https://hypothesis.works/articles/what-is-property-based-testing/",
      "label": "fuzzy",
      "type": "a"
    },
    {
      "href": "https://stackoverflow.com/questions/520064/what-are-the-differences-between-unit-tests-integration-tests-smoke-tests-and",
      "label": "What are the Differences between Unit Tests, Integration Tests, Smoke Tests, etc?",
      "type": "a"
    },
    {
      "href": "https://www.slideshare.net/hkarau/effective-testing-for-spark-programs-strata-ny-2015",
      "label": "Effective Testing for Spark Programs",
      "type": "a"
    },
    {
      "href": "https://www.slideshare.net/mathieu-bastian/the-mechanics-of-testing-large-data-pipelines-qcon-london-2016",
      "label": "The Mechanics of Testing Large Data Pipelines",
      "type": "a"
    },
    {
      "href": "https://www.slideshare.net/lallea/test-strategies-for-data-processing-pipelines-67244458",
      "label": "Test Strategies for Data Processing Pipelines",
      "type": "a"
    },
    {
      "href": "https://www.slideshare.net/lallea/engineering-data-quality",
      "label": "Engineering Data Quality",
      "type": "a"
    },
    {
      "href": "https://towardsdatascience.com/testing-features-with-pytest-82765a13a0e7",
      "label": "Testing Features with PyTest",
      "type": "a"
    },
    {
      "href": "https://www.freecodecamp.org/news/intro-to-property-based-testing-in-python-6321e0c2f8b/",
      "label": "Intro to Property-based Testing in Python",
      "type": "a"
    },
    {
      "href": "https://github.com/microsoft/recommenders/tree/main/tests",
      "label": "Testing in Microsoft Recommenders",
      "type": "a"
    },
    {
      "href": "https://www.linkedin.com/in/david-said-martinez-3836a72b/",
      "label": "David Said",
      "type": "a"
    },
    {
      "href": "https://twitter.com/jim_dowling",
      "label": "Jim Dowling",
      "type": "a"
    },
    {
      "href": "https://twitter.com/sh_reya",
      "label": "Shreya Shankar",
      "type": "a"
    },
    {
      "href": "https://www.linkedin.com/in/mikeyplam/",
      "label": "Mike Lam",
      "type": "a"
    },
    {
      "href": "https://twitter.com/lalleal",
      "label": "Lars Albertsson",
      "type": "a"
    },
    {
      "href": "https://unsplash.com/@samloyd",
      "label": "Sam Loyd",
      "type": "a"
    },
    {
      "href": "https://unsplash.com/s/photos/pipes",
      "label": "Unsplash",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/writing/simplicity/",
      "label": "« Simplicity is An Advantage but Sadly Complexity Sells Better",
      "type": "a"
    },
    {
      "href": "https://eugeneyan.com/speaking/recsys2022-keynote/",
      "label": "RecSys 2022 Keynote - Is the Juice Worth the Squeeze? »",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:52:20Z",
  "source_url": "https://eugeneyan.com/writing/testing-pipelines/",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 27766,
  "expected_body_images": 7,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/b04dabee6e8d/pages/source.html",
  "local_html": "attachments/b04dabee6e8d/pages/article.html",
  "body_base_path": "attachments/b04dabee6e8d/pages/",
  "body_markdown_path": "attachments/b04dabee6e8d/pages/article.md"
}