diff --git a/.gitignore b/.gitignore
index 7773828..c5318dc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
-dist/
\ No newline at end of file
+dist/
+**/__pycache__/
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 3350f84..9d95591 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,6 +35,7 @@ classifiers = [
 # path = "src/fy_polyfit/polyfit.py"
 source = "vcs"
 
+
 [project.scripts]
 polyfit = "fy_polyfit.polyfit:main"
 
diff --git a/src/fy_polyfit/__pycache__/polyfit.cpython-311.pyc b/src/fy_polyfit/__pycache__/polyfit.cpython-311.pyc
deleted file mode 100644
index 62ae282..0000000
Binary files a/src/fy_polyfit/__pycache__/polyfit.cpython-311.pyc and /dev/null differ
diff --git a/src/fy_polyfit/polyfit.py b/src/fy_polyfit/polyfit.py
index 2e4c92e..db8aa72 100644
--- a/src/fy_polyfit/polyfit.py
+++ b/src/fy_polyfit/polyfit.py
@@ -26,6 +26,12 @@ def  cal_rmse(y_true, y_pred):
     return np.sqrt(np.mean((y_pred - y_true) **2))
 
 
+def PolyfitLinear(x, y):
+     return np.polyfit(x, y, 1)
+
+def PolyfitQuadratic(x, y):
+     return np.polyfit(x, y, 2)
+
 def main():
     parser = argparse.ArgumentParser(description='从指定文件中读取数据,默认为data.txt, 无头部,空格分隔,两列')
     parser.add_argument('filename', nargs='?', default='data.txt', 
@@ -108,8 +114,12 @@ def main():
         ax[1].text(0.5, 0.8, formula2,color='blue',fontsize=16, transform=ax[1].transAxes, verticalalignment='top', horizontalalignment='center')
 
         plt.tight_layout()
+ 
+        # plt.draw()
+        # 保存图片到images文件夹下
+        plt.savefig('../../images/fitting.png')
         plt.show()
-
+    
     except FileNotFoundError as e:
         print(e)
     
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_polyfit.py b/tests/test_polyfit.py
new file mode 100644
index 0000000..bc3744b
--- /dev/null
+++ b/tests/test_polyfit.py
@@ -0,0 +1,36 @@
+import sys, os
+import numpy as np
+import pytest
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+from fy_polyfit.polyfit import PolyfitLinear, PolyfitQuadratic
+
+def test_polyfit():
+    x = np.array([1, 2, 3, 4, 5])
+    y = 2 * x + 1 + np.random.uniform(-0.5, 0.5, size=x.shape)
+    
+    # 调用 PolyfitLinear
+    coefficients = PolyfitLinear(x, y)
+    
+    # 期望的系数
+    expected_coefficients = np.polyfit(x, y, 1)
+    
+    # 断言系数接近期望值
+    np.testing.assert_almost_equal(coefficients, expected_coefficients, decimal=2)
+
+def test_polyfit_quadratic():
+    # 示例数据
+    x = np.array([1, 2, 3, 4, 5])
+    y = 3 * x**2 + 2 * x + 1 + np.random.uniform(-1, 1, size=x.shape)
+    
+    # 调用 PolyfitQuadratic
+    coefficients = PolyfitQuadratic(x, y)
+    
+    # 期望的系数
+    expected_coefficients = np.polyfit(x, y, 2)
+    
+    # 断言系数接近期望值
+    np.testing.assert_almost_equal(coefficients, expected_coefficients, decimal=2)
+
+
+if __name__ == "__main__":
+    pytest.main()
\ No newline at end of file