Testing Master Mind

import random


def main():
    hidden = [str(x) for x in random.sample(range(1, 7), 4)]
    #print(hidden)
    
    while(True):
        print("Please enter 4 digits")
        guess = list(input())
        if len(guess)!=4:
            continue    
        res = ""
        for h, g in zip(hidden, guess):
           #print(h, g)
           if h == g:
              res += "b"
           elif g in hidden:
              res += "w"
           #print(res)
        if res=='bbbb':
           print("Congrats!") 
           break   
        print(''.join(sorted(res)))


if __name__ == "__main__":
    main() 

import master_mind as mm
import random

def test_mm():
    random.sample = lambda a, b: [1,2,3,4]
    input_values = ['1234']
    output = []

    def mock_input():
       #output.append(s)
       return input_values.pop(0)
    mm.input = mock_input
    mm.print = lambda *s : output.append(s)

    mm.main()

    assert output == [
        ("Please enter 4 digits",),
        ('Congrats!',),
    ] 


def test_wrong():
    random.sample = lambda a, b: [1,2,3,4]
    input_values = ['1235', '1234']
    output = []

    def mock_input():
       #output.append(s)
       return input_values.pop(0)
    mm.input = mock_input
    mm.print = lambda *s : output.append(s)

    mm.main()

    assert output == [
        ("Please enter 4 digits",),
        ("bbb",),
        ("Please enter 4 digits",),
        ('Congrats!',),
    ]