In A/B testing, assigning users to control and variant groups in a balanced and unbiased way is essential for reliable results. If the groups differ systematically in user characteristics - such as age, location, or behavior - any observed effect might be due to these differences rather than the change being tested. **Randomization** helps ensure that each user has an equal chance of being assigned to any group, which minimizes bias and helps create comparable groups.

To check for balance, you can compare the distributions of user characteristics between the groups after assignment. Ideally, both groups should have similar averages and proportions for important variables. If they are not balanced, it may indicate a flaw in the randomization process or the presence of confounding variables. Careful assignment and balance checking help ensure the validity of your A/B test conclusions.


import unittest
import user_code
import ast
import re   
import importlib
import csv
import unittest
import pandas as pd
import numpy as np
import importlib
import sys

class TestTask(unittest.TestCase):
    def setUp(self):
        self.users = pd.DataFrame({
            'user_id': [101, 102, 103, 104, 105, 106, 107, 108],
            'age': [23, 45, 31, 29, 40, 22, 30, 27],
            'location': ['NY', 'CA', 'TX', 'NY', 'CA', 'TX', 'NY', 'CA']
        })

    def test_group_column_exists_and_valid(self):
        import user_code
        importlib.reload(user_code)
        df = user_code.assign_groups(self.users)
        _dynamic_test(
            self,
            'group' in df.columns,
            "Returned DataFrame includes a 'group' column.",
            "Returned DataFrame does not have a 'group' column."
        )
        valid = set(df['group'].unique()) <= {'control', 'variant'}
        _dynamic_test(
            self,
            valid,
            "All group assignments are either 'control' or 'variant'.",
            f"Group assignments must be only 'control' or 'variant', got: {set(df['group'].unique())}"
        )

    def test_each_user_one_group(self):
        import user_code
        importlib.reload(user_code)
        df = user_code.assign_groups(self.users)
        counts = df.groupby('user_id')['group'].nunique().values
        _dynamic_test(
            self,
            all(c == 1 for c in counts),
            "Each user is assigned to exactly one group.",
            f"Some users assigned to multiple groups: group counts: {counts}"
        )

    def test_assignment_is_random_and_balanced(self):
        import user_code
        # Run multiple times to check for randomness
        group_counts = []
        for _ in range(20):
            importlib.reload(user_code)
            df = user_code.assign_groups(self.users)
            counts = df['group'].value_counts()
            group_counts.append(counts)
        # Check that both groups appear in all runs
        for counts in group_counts:
            _dynamic_test(
                self,
                'control' in counts and 'variant' in counts,
                "Both groups are assigned in each run.",
                f"One of the groups is missing in assignment: {counts.to_dict()}"
            )
        # Check that distribution is approximately even
        diffs = [abs(counts.get('control',0) - counts.get('variant',0)) for counts in group_counts]
        _dynamic_test(
            self,
            all(d <= 4 for d in diffs),
            "Assignment is approximately balanced for 8 users.",
            f"Assignment is not balanced enough: diffs={diffs}"
        )

    def test_original_dataframe_not_modified(self):
        import user_code
        importlib.reload(user_code)
        orig = self.users.copy()
        _ = user_code.assign_groups(self.users)
        _dynamic_test(
            self,
            self.users.equals(orig),
            "Original DataFrame is not modified.",
            "Original DataFrame was changed after function call."
        )

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\\s{2,}", " ", text)
    text = re.sub(r"\\s*([,:?])\\s*", r"\\1 ", text)
    return text.strip()

def change_var(code: str, var_name: str, value: str) -> str:
    tree = ast.parse(code)
    lines = code.splitlines()
    changed = False
    # Collect all assignment nodes to modify
    assign_nodes = [
        (i, node)
        for i, node in enumerate(tree.body)
        if isinstance(node, ast.Assign)
        and any(isinstance(target, ast.Name) and target.id == var_name for target in node.targets)
    ]

    # If nothing to change, return unmodified code
    if not assign_nodes:
        return code

    # Perform replacements for all matching assignments (from last to first to not break line offsets)
    for i, node in reversed(assign_nodes):
        start_line = node.lineno - 1
        line = lines[start_line]
        indent = ' ' * (len(line) - len(line.lstrip()))
        lines[start_line] = f"{indent}{var_name} = {value}"
        next_line = len(lines)
        for next_node in tree.body[i+1:]:
            if hasattr(next_node, 'lineno'):
                next_line = next_node.lineno - 1
                break
        if next_line > start_line + 1:
            lines[start_line+1:next_line] = []
        changed = True

    return '\\n'.join(lines) if changed else code

if __name__ == "__main__":
    unittest.main()


Challenge: Assigning Users to Groups

Solution